ic

apps/ChatEgresos/api/agent/main.py

@@ -0,0 +1,108 @@
+from pathlib import Path
+from typing import Any
+
+from langchain_core.messages import AIMessageChunk
+from pydantic import BaseModel, Field
+from langchain_azure_ai.chat_models import AzureAIChatCompletionsModel
+from langchain_azure_ai.embeddings import AzureAIEmbeddingsModel
+
+from banortegpt.vector.qdrant import AsyncQdrant
+
+from api import context
+from api.config import config
+
+parent = Path(__file__).parent
+SYSTEM_PROMPT = (parent / "system_prompt.md").read_text()
+
+AZURE_AI_URI = "https://eastus2.api.cognitive.microsoft.com"
+
+class get_information(BaseModel):
+    """Search a private repository for information."""
+
+    question: str = Field(..., description="The user question")
+
+class Agent:
+    system_prompt = SYSTEM_PROMPT
+    generation_config = {
+        "temperature": config.model_temperature,
+    }
+    embedding_model = config.embedding_model
+    message_limit = config.message_limit
+    index = config.vector_index
+    limit = config.search_limit
+
+    search = AsyncQdrant.from_config(config)
+    llm = AzureAIChatCompletionsModel(
+        endpoint=f"{AZURE_AI_URI}/openai/deployments/{config.model}",
+        credential=config.openai_api_key,
+    ).bind_tools([get_information])
+    embedder = AzureAIEmbeddingsModel(
+        endpoint=f"{AZURE_AI_URI}/openai/deployments/{config.embedding_model}",
+        credential=config.openai_api_key,
+    )
+  
+    def __init__(self) -> None:
+        self.tool_map = {
+            "get_information": self.get_information
+        }
+
+    def build_response(self, payloads, fallback):
+        template = "<FAQ {index}>\n\n{content}\n\n</FAQ {index}>"
+        
+        filled_templates = [
+            template.format(index=idx, content=payload["content"])
+            for idx, payload in enumerate(payloads)
+        ]
+        filled_templates.append(f"<FALLBACK>\n{fallback}\n</FALLBACK>")
+        
+        return "\n".join(filled_templates)
+
+    async def get_information(self, question: str):
+        embedding = await self.embedder.aembed_query(question)
+        
+        payloads = await self.search.semantic_search(
+            embedding=embedding,
+            collection=self.index,
+            limit=self.limit,
+        )
+        
+        fallback_messages = {}
+        images = []
+        for idx, payload in enumerate(payloads):
+            fallback_message = payload.get("fallback_message", "None")
+            fallback_messages[fallback_message] = fallback_messages.get(fallback_message, 0) + 1
+
+            # Solo extraer imágenes del primer payload
+            if idx == 0 and "images" in payload:    
+                images.extend(payload["images"])
+
+        fallback = max(fallback_messages, key=fallback_messages.get)  # type: ignore
+        
+        response = self.build_response(payloads, fallback)
+        return str(response), images[:3]  # Limitar a 3 imágenes máximo
+
+    def _generation_config_overwrite(self, overwrites: dict | None) -> dict[str, Any]:
+        if not overwrites:
+            return self.generation_config.copy()
+        return {**self.generation_config, **overwrites}
+
+    async def stream(self, history, overwrites: dict | None = None):
+        generation_config = self._generation_config_overwrite(overwrites)
+
+        async for delta in self.llm.astream(input=history, **generation_config):
+            assert isinstance(delta, AIMessageChunk)
+            if call := delta.tool_call_chunks:
+                if tool_id := call[0].get("id"):
+                    context.tool_id.set(tool_id)
+                if name := call[0].get("name"):
+                    context.tool_name.set(name)
+                if args := call[0].get("args"):
+                    context.tool_buffer.set(context.tool_buffer.get() + args)
+            elif delta.content:
+                assert isinstance(delta.content, str)
+                context.buffer.set(context.buffer.get() + delta.content)
+                yield delta.content
+
+    async def generate(self, history, overwrites: dict | None = None):
+        generation_config = self._generation_config_overwrite(overwrites)
+        return await self.llm.ainvoke(input=history, **generation_config)