118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
import asyncio
|
|
from langchain.agents import create_react_agent, AgentExecutor
|
|
from langchain.chat_models import init_chat_model
|
|
from langchain_core.prompts import PromptTemplate
|
|
from langchain_mcp_adapters.client import MultiServerMCPClient
|
|
from langchain_core.tools import tool
|
|
from langchain.memory import MongoDBChatMessageHistory, ConversationBufferWindowMemory
|
|
|
|
from agents.agent_templates import AgentTemplate
|
|
from agents.output_parser.MCPReactParser import MCPReactParser
|
|
|
|
|
|
@tool
|
|
def get_daily_string(salt: str):
|
|
"""
|
|
:param salt: a random string
|
|
:return:get today's string
|
|
"""
|
|
return "assss"
|
|
|
|
|
|
async def main():
|
|
|
|
temp1 = AgentTemplate()
|
|
temp1.set_model("gpt-4.1-nano", "openai", "")
|
|
temp1.set_mcp_server("terminal_tools", {
|
|
'transport': 'sse',
|
|
'headers': {
|
|
"X-Api-Key": "bd303c2f2a515b4ce70c1f07ee85530c2"
|
|
},
|
|
'url': 'http://localhost:5050/sse'
|
|
})
|
|
|
|
|
|
mcp_client = MultiServerMCPClient({
|
|
'terminal_tool':{
|
|
'transport': 'sse',
|
|
'headers': {
|
|
"X-Api-Key": "bd303c2f2a515b4ce70c1f07ee85530c2"
|
|
},
|
|
'url': 'http://localhost:5050/sse'
|
|
}
|
|
})
|
|
|
|
tools = await mcp_client.get_tools()
|
|
xtools = []
|
|
for t in tools:
|
|
if t.name == 'CreateTerminalSession':
|
|
xtools.append(t)
|
|
xtools.append(get_daily_string)
|
|
|
|
print(xtools)
|
|
prompt = PromptTemplate.from_template("""
|
|
|
|
{system}
|
|
|
|
You have access to the following tools:
|
|
|
|
{tools}
|
|
|
|
Use the following format:
|
|
|
|
Question: the question you must answer
|
|
If you want to use tools:
|
|
Thought: always reason what to do
|
|
Action: the action to take, must be one of [{tool_names}]
|
|
Action Input: a **JSON object** with named arguments required by the tool
|
|
Observation: the result of the action
|
|
If no tool is needed:
|
|
Thought: what you are thinking
|
|
... (this Thought/Action/... can repeat N times)
|
|
Final Answer: the final answer to the original question
|
|
|
|
Here is the conversation history:
|
|
{chat_history}
|
|
|
|
User message: {user}
|
|
{agent_scratchpad}
|
|
""")
|
|
op = MCPReactParser()
|
|
agent = create_react_agent(model_openai, xtools, prompt, output_parser=op)
|
|
message_history = MongoDBChatMessageHistory(
|
|
connection_string=MONGO_CONNECTION_STRING,
|
|
session_id=USER_SESSION_ID,
|
|
database_name=MONGO_DB_NAME,
|
|
collection_name=MONGO_COLLECTION_NAME
|
|
)
|
|
mongodb_memory = ConversationBufferWindowMemory(
|
|
chat_memory=message_history,
|
|
memory_key=MEMORY_KEY,
|
|
k=10,
|
|
return_messages=False,
|
|
)
|
|
agent_executor = AgentExecutor(
|
|
agent=agent,
|
|
tools=xtools,
|
|
memory=mongodb_memory,
|
|
handle_parsing_errors=True,
|
|
verbose=True,
|
|
|
|
)
|
|
response = await agent_executor.ainvoke({
|
|
'system': 'you are a helpful assistant',
|
|
'user': 'hello, please create a terminal session with label \'x\' and show me the session id',
|
|
})
|
|
print(response)
|
|
r2 = await agent_executor.ainvoke({
|
|
'system': 'you are a helpful assistant',
|
|
'user': 'hello, please show me today\'s daily string'
|
|
})
|
|
print(r2)
|
|
#print(s1)
|
|
#print(s2)
|
|
#print(s3)
|
|
#print(s4)
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main()) |