ic

apps/pyme/api/services/__init__.py

@@ -0,0 +1,8 @@
+from banortegpt.database.mongo_memory.crud import create_conversation
+
+from .stream_response import stream
+
+__all__ = [
+    "stream",
+    "create_conversation",
+]

apps/pyme/api/services/stream_response.py

@@ -0,0 +1,89 @@
+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="")