ic

apps/ChatEgresos/api/services/stream_response.py

@@ -0,0 +1,86 @@
+import json
+from enum import StrEnum
+from typing import TypeAlias
+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
+
+
+class ResponseChunk(BaseModel):
+    type: ChunkType
+    content: ContentType | list[ContentType] | None
+    images: list[str] | None = None  # Nuevo campo para imágenes
+
+
+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)
+
+    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, images = await agent.tool_map[tool_name](**buffer_dict)
+        
+        # Enviar imágenes si existen
+        if images:
+            yield ResponseChunk(type=ChunkType.IMAGE, content=images)
+
+        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="")