forked from innovacion/Mayacontigo
89 lines
2.4 KiB
Python
89 lines
2.4 KiB
Python
import json
|
|
from enum import StrEnum
|
|
from typing import TypeAlias, Any
|
|
from uuid import UUID
|
|
|
|
from pydantic import BaseModel
|
|
|
|
import api.context as ctx
|
|
from api.agent import Agent
|
|
from banortegpt.database.mongo_memory import crud
|
|
|
|
|
|
class ChunkType(StrEnum):
|
|
START = "start"
|
|
TEXT = "text"
|
|
REFERENCE = "reference"
|
|
IMAGE = "image"
|
|
TOOL = "tool"
|
|
END = "end"
|
|
ERROR = "error"
|
|
|
|
|
|
ContentType: TypeAlias = str | int | dict | Any
|
|
|
|
|
|
class ResponseChunk(BaseModel):
|
|
type: ChunkType
|
|
content: ContentType | list[ContentType] | None
|
|
|
|
|
|
async def stream(agent: Agent, prompt: str, conversation_id: UUID):
|
|
yield ResponseChunk(type=ChunkType.START, content="")
|
|
|
|
conversation = await crud.get_conversation(conversation_id)
|
|
|
|
if conversation is None:
|
|
raise ValueError("Conversation not found")
|
|
|
|
conversation.add(role="user", content=prompt)
|
|
|
|
history = conversation.to_openai_format(agent.message_limit, langchain_compat=True)
|
|
|
|
async for content in agent.stream(history):
|
|
yield ResponseChunk(type=ChunkType.TEXT, content=content)
|
|
|
|
|
|
metadatas = getattr(agent, 'last_metadatas', [])
|
|
|
|
|
|
if metadatas:
|
|
yield ResponseChunk(type=ChunkType.REFERENCE, content=metadatas)
|
|
|
|
if (tool_id := ctx.tool_id.get()) is not None:
|
|
tool_buffer = ctx.tool_buffer.get()
|
|
assert tool_buffer is not None
|
|
|
|
tool_name = ctx.tool_name.get()
|
|
assert tool_name is not None
|
|
|
|
yield ResponseChunk(type=ChunkType.TOOL, content=None)
|
|
|
|
buffer_dict = json.loads(tool_buffer)
|
|
|
|
result = await agent.tool_map[tool_name](**buffer_dict)
|
|
|
|
conversation.add(
|
|
role="assistant",
|
|
tool_calls=[
|
|
{
|
|
"id": tool_id,
|
|
"type": "function",
|
|
"function": {
|
|
"name": tool_name,
|
|
"arguments": tool_buffer,
|
|
},
|
|
}
|
|
],
|
|
)
|
|
conversation.add(role="tool", content=result, tool_call_id=tool_id)
|
|
|
|
history = conversation.to_openai_format(agent.message_limit, langchain_compat=True)
|
|
async for content in agent.stream(history, {"tools": None}):
|
|
yield ResponseChunk(type=ChunkType.TEXT, content=content)
|
|
|
|
conversation.add(role="assistant", content=ctx.buffer.get())
|
|
|
|
await conversation.replace()
|
|
|
|
yield ResponseChunk(type=ChunkType.END, content="") |