Spaces:
Build error
Build error
File size: 4,069 Bytes
79a4fbe e0ba3a7 79a4fbe e0ba3a7 79a4fbe e0ba3a7 79a4fbe e0ba3a7 79a4fbe e0ba3a7 79a4fbe e0ba3a7 79a4fbe e0ba3a7 79a4fbe e0ba3a7 79a4fbe e0ba3a7 79a4fbe |
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 |
import json
from mcp import ClientSession
import anthropic
import chainlit as cl
anthropic_client = anthropic.AsyncAnthropic()
SYSTEM = "You are a helpful stock trader that help me identify opportunities and risks in the stock market. " \
"Provide brief and concise reports of current market conditions, including key indicators and trends. " \
"Use the tools provided to gather data and insights. " \
"If you need to use a tool, please do so and provide the results in your response. " \
"You can also ask clarifying questions if needed. " \
"You try to maximize the profit of the user. "
def flatten(xss):
return [x for xs in xss for x in xs]
@cl.on_mcp_connect
async def on_mcp(connection, session: ClientSession):
result = await session.list_tools()
tools = [{
"name": t.name,
"description": t.description,
"input_schema": t.inputSchema,
} for t in result.tools]
mcp_tools = cl.user_session.get("mcp_tools", {})
mcp_tools[connection.name] = tools
cl.user_session.set("mcp_tools", mcp_tools)
@cl.step(type="tool")
async def call_tool(tool_use):
tool_name = tool_use.name
tool_input = tool_use.input
current_step = cl.context.current_step
current_step.name = tool_name
# Identify which mcp is used
mcp_tools = cl.user_session.get("mcp_tools", {})
mcp_name = None
for connection_name, tools in mcp_tools.items():
if any(tool.get("name") == tool_name for tool in tools):
mcp_name = connection_name
break
if not mcp_name:
current_step.output = json.dumps({"error": f"Tool {tool_name} not found in any MCP connection"})
return current_step.output
mcp_session, _ = cl.context.session.mcp_sessions.get(mcp_name)
if not mcp_session:
current_step.output = json.dumps({"error": f"MCP {mcp_name} not found in any MCP connection"})
return current_step.output
try:
current_step.output = await mcp_session.call_tool(tool_name, tool_input)
except Exception as e:
current_step.output = json.dumps({"error": str(e)})
return current_step.output
async def call_claude(chat_messages):
msg = cl.Message(content="")
mcp_tools = cl.user_session.get("mcp_tools", {})
# Flatten the tools from all MCP connections
tools = flatten([tools for _, tools in mcp_tools.items()])
async with anthropic_client.messages.stream(
system=SYSTEM,
max_tokens=1024,
messages=chat_messages,
tools=tools,
model="claude-3-5-sonnet-20240620",
) as stream:
async for text in stream.text_stream:
await msg.stream_token(text)
await msg.send()
response = await stream.get_final_message()
return response
@cl.on_chat_start
async def start_chat():
cl.user_session.set("chat_messages", [])
@cl.on_message
async def on_message(msg: cl.Message):
chat_messages = cl.user_session.get("chat_messages")
chat_messages.append({"role": "user", "content": msg.content})
response = await call_claude(chat_messages)
while response.stop_reason == "tool_use":
tool_use = next(block for block in response.content if block.type == "tool_use")
tool_result = await call_tool(tool_use)
messages = [
{"role": "assistant", "content": response.content},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": tool_use.id,
"content": str(tool_result),
}
],
},
]
chat_messages.extend(messages)
response = await call_claude(chat_messages)
final_response = next(
(block.text for block in response.content if hasattr(block, "text")),
None,
)
chat_messages = cl.user_session.get("chat_messages")
chat_messages.append({"role": "assistant", "content": final_response}) |