File size: 4,652 Bytes
f647629
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7b2da21
f647629
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import requests
import os
from typing import Any, Dict

# from wandbot.wandbot_tool import query_wandbot_api, WANDBOT_TOOL_DESCRIPTION

WANDBOT_TOOL_DESCRIPTION = """Query the Weights & Biases support bot api for help with questions about the
Weights & Biases platform and how to use W&B Models and W&B Weave.

W&B features mentioned could include:
- Experiment tracking with Runs and Sweeps
- Model management with Models
- Model management and Data versioning with Artifacts and Registry
- Collaboration with Teams, Organizations and Reports
- Visualization with Tables and Charts
- Tracing and logging with Weave
- Evaluation and Scorers with Weave Evaluations
- Weave Datasets

FYI: The Weigths & Biases platform is owned by Coreweave. If there are queries related to W&B, wandb \
or weave and Coreweave, they might be related to W&B products or features that leverage Coreweave's \
GPU or compute infrastructure.

Parameters
----------
question : str
    Users question about a Weights & Biases product or feature

Returns
-------
str
    newer to the user's question
"""


def query_wandbot_api(question: str) -> Dict[str, Any]:
    wandbot_base_url = os.getenv("WANDBOT_BASE_URL", "https://weightsandbiases-wandbot--wandbot-api-wandbotapi-serve.modal.run")
    QUERY_ENDPOINT = f"{wandbot_base_url}/chat/query"
    STATUS_ENDPOINT = f"{wandbot_base_url}/status"
    QUERY_TIMEOUT_SECONDS = 40
    STATUS_TIMEOUT_SECONDS = 20

    try:
        status_response = requests.get(
            STATUS_ENDPOINT,
            headers={"Accept": "application/json"},
            timeout=STATUS_TIMEOUT_SECONDS,
        )

        # Check HTTP status code
        status_response.raise_for_status()

        # Try to parse JSON, handle potential parsing errors
        try:
            status_result = status_response.json()
        except ValueError:
            return {
                "answer": "Error: Unable to parse response from support bot.",
                "sources": [],
            }

        # Validate expected response structure
        if "initialized" not in status_result:
            return {
                "answer": "Error: Received unexpected response format from support bot.",
                "sources": [],
            }

        if status_result["initialized"]:
            try:
                response = requests.post(
                    QUERY_ENDPOINT,
                    headers={"Content-Type": "application/json"},
                    json={
                        "question": question,
                        "application": "wandb_mcp_server",
                    },
                    timeout=QUERY_TIMEOUT_SECONDS,
                )

                # Check HTTP status code
                response.raise_for_status()

                # Try to parse JSON, handle potential parsing errors
                try:
                    result = response.json()
                except ValueError:
                    return {
                        "answer": "Error: Unable to parse response data from support bot.",
                        "sources": [],
                    }

                # Validate expected response structure
                if "answer" not in result or "sources" not in result:
                    return {
                        "answer": "Error: Received incomplete response from support bot.",
                        "sources": [],
                    }

                # Ensure sources is a list
                sources = (
                    result["sources"]
                    if isinstance(result["sources"], list)
                    else [result["sources"]]
                )
                return {"answer": result["answer"], "sources": sources}

            except requests.Timeout:
                return {
                    "answer": "Error: Support bot request timed out. Please try again later.",
                    "sources": [],
                }
            except requests.RequestException as e:
                return {
                    "answer": f"Error connecting to support bot: {str(e)}",
                    "sources": [],
                }
        else:
            return {
                "answer": "The support bot is appears to be offline. Please try again later.",
                "sources": [],
            }

    except requests.Timeout:
        return {
            "answer": "Error: Support bot status check timed out. Please try again later.",
            "sources": [],
        }
    except requests.RequestException as e:
        return {
            "answer": f"Error connecting to support bot: {str(e)}",
            "sources": [],
        }