57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import asyncio
|
|
import traceback
|
|
|
|
from lambda_agent_api.agent_api import AgentApi, AgentBusyException
|
|
from lambda_agent_api.server import (
|
|
MsgEventTextChunk,
|
|
MsgEventToolCallChunk,
|
|
MsgEventToolResult,
|
|
MsgEventSendFile,
|
|
)
|
|
|
|
|
|
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)
|
|
case MsgEventSendFile():
|
|
print(f"\n[SEND FILE] {chunk.path}\n", end="", flush=True)
|
|
|
|
print("\n")
|
|
except KeyboardInterrupt:
|
|
break
|
|
except:
|
|
traceback.print_exc()
|
|
|
|
await api.close()
|
|
|
|
|
|
asyncio.run(main())
|