File size: 2,869 Bytes
2329760 |
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 |
import os
import requests
from langchain_mcp_adapters.tool import MCPTool
gaia_system_prompt = """
You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
"""
class BasicAgent:
def __init__(self, model_name="deepseek-ai/deepseek-v3.1"):
print("BasicAgent initialized.")
# MCP tool configuration
self.mcp_tool = MCPTool(
tool_name="generic_tool", # replace with actual tool name exposed by MCP
server_url="http://localhost:8080",
)
# NVIDIA NIM API configuration
self.model_name = model_name
self.api_key = os.getenv("NVIDIA_API_KEY")
self.api_base = "https://integrate.api.nvidia.com/v1"
def call_nim_api(self, user_input: str) -> str:
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
payload = {
"model": self.model_name,
"messages": [
{"role": "system", "content": gaia_system_prompt},
{"role": "user", "content": user_input}
],
"temperature": 0.7
}
response = requests.post(
f"{self.api_base}/chat/completions",
headers=headers,
json=payload
)
try:
return response.json()["choices"][0]["message"]["content"]
except Exception as e:
print("Error calling NIM API:", e)
return "NIM API call failed."
def __call__(self, question: str) -> str:
print(f"Agent received input (first 50 chars): {question[:50]}...")
# Call NVIDIA NIM API
nim_output = self.call_nim_api(question)
print(f"NIM response: {nim_output[:100]}...")
# Optionally use MCP tool based on input
if "scrape" in question.lower():
mcp_result = self.mcp_tool.run({
"url": "https://example.com",
"selectors": {
"title": ".title",
"price": ".price"
}
})
print("MCP result:", mcp_result)
return f"NIM: {nim_output}\n\nMCP: {mcp_result}"
return nim_output
|