forked from trpc-group/trpc-agent-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathrun_agent.py
More file actions
132 lines (107 loc) · 4.49 KB
/
Copy pathrun_agent.py
File metadata and controls
132 lines (107 loc) · 4.49 KB
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
#!/usr/bin/env python3
# Tencent is pleased to support the open source community by making tRPC-Agent-Python available.
#
# Copyright (C) 2026 Tencent. All rights reserved.
#
# tRPC-Agent-Python is licensed under Apache-2.0.
"""Run the weather query agent demo"""
import asyncio
import sys
from pathlib import Path
from dotenv import load_dotenv
from trpc_agent_sdk.context import AgentContext
from trpc_agent_sdk.memory import MemoryServiceConfig
from trpc_agent_sdk.memory.mempalace_memory_service import MempalaceMemoryService
from trpc_agent_sdk.runners import Runner
from trpc_agent_sdk.sessions import InMemorySessionService
from trpc_agent_sdk.types import Content
from trpc_agent_sdk.types import Part
# Load environment variables from the .env file
load_dotenv()
sys.path.append(str(Path(__file__).parent))
def _truncate_tool_response(response, max_length: int = 256) -> str:
"""Limit verbose tool responses in demo output."""
text = str(response)
if len(text) <= max_length:
return text
return text[:max_length]
def create_memory_service():
"""Create session service"""
memory_service_config = MemoryServiceConfig(
ttl=MemoryServiceConfig.create_ttl_config(enable=True, ttl_seconds=20, cleanup_interval_seconds=20),
enabled=True,
)
memory_service = MempalaceMemoryService(
memory_service_config=memory_service_config,
wing="trpc-agent",
room="conversations",
)
return memory_service
async def run_weather_agent(memory_service: MempalaceMemoryService):
"""Run the weather query agent demo"""
app_name = "weather_agent_demo"
from agent.agent import root_agent
session_service = InMemorySessionService()
runner = Runner(app_name=app_name, agent=root_agent, session_service=session_service, memory_service=memory_service)
user_id = "mempalace_memory_user"
current_session_id = "mempalace_memory_session"
# Demo query list
demo_queries = [
"Do you remember my name?",
"Do you remember my favorite color?",
"What is the weather like in paris?",
"Hello! My name is Alice. Please remember my name.",
"Now, do you still remember my name?",
"Hello! My favorite color is blue. Please remember my favorite color.",
"Now, do you still remember my favorite color?",
]
for index, query in enumerate(demo_queries):
# Use a new session for each query
user_content = Content(parts=[Part.from_text(text=query)])
print("🤖 Assistant: ", end="", flush=True)
agent_context = AgentContext()
session_id = f"{current_session_id}_{index}"
# set_mem0_filters(agent_context, {"session_id": session_id})
async for event in runner.run_async(agent_context=agent_context,
user_id=user_id,
session_id=session_id,
new_message=user_content):
# Check if event.content exists
if not event.content or not event.content.parts:
continue
if event.partial:
for part in event.content.parts:
if part.text:
print(part.text, end="", flush=True)
continue
for part in event.content.parts:
# Skip the reasoning part; the output is already generated when partial=True
if part.thought:
continue
if part.function_call:
print(f"\n🔧 [Invoke Tool: {part.function_call.name}({part.function_call.args})]")
elif part.function_response:
print(f"📊 [Tool Result: {_truncate_tool_response(part.function_response.response)}]")
# Uncomment to get the full text output of the LLM
# elif part.text:
# print(f"\n✅ {part.text}")
print("\n" + "-" * 40)
async def main():
memory_service = create_memory_service()
print("=" * 60)
print("First run")
print("=" * 60)
await run_weather_agent(memory_service)
await asyncio.sleep(2)
print("=" * 60)
print("Second run")
print("=" * 60)
await run_weather_agent(memory_service)
await asyncio.sleep(30)
print("=" * 60)
print("Third run")
print("=" * 60)
await run_weather_agent(memory_service)
# await memory_service.delete_memory(wing="trpc-agent", room="conversations")
if __name__ == "__main__":
asyncio.run(main())