agent_api/tests/manual.py

47 lines
No EOL
1.5 KiB
Python

import asyncio
import traceback
from lambda_agent_api.agent_api import AgentApi, AgentBusyException
from lambda_agent_api.server import MsgEventTextChunk, MsgEventToolCallChunk, MsgEventToolResult
async def main():
chat_id = input("Chat id: ") or 0
api = AgentApi("agent-1", "ws://localhost:8000/", chat_id=chat_id)
try:
await api.connect()
except AgentBusyException:
print(f"Чат {chat_id} занят другим клиентом")
return
while True:
try:
prompt = await asyncio.get_event_loop().run_in_executor(None, input, ">>> ")
print("Agent: ", end="")
is_tool = False
async for chunk in api.send_message(prompt):
match chunk:
case MsgEventTextChunk():
is_tool = False
print(chunk.text, end="", flush=True)
case MsgEventToolCallChunk():
if not is_tool:
print(f"\n\n### TOOL CALL: ({chunk.tool_name}) ", end="", flush=True)
is_tool = True
print(chunk.args_chunk, end="", flush=True)
case MsgEventToolResult():
is_tool = False
print(f"\nResult: {chunk.result}\n\n", end="", flush=True)
print("\n")
except KeyboardInterrupt:
break
except:
traceback.print_exc()
await api.close()
asyncio.run(main())