ssss

my_tools.py

@@ -47,66 +47,132 @@ ChatMessage.message = property(lambda self: self)
 
 # ---------- GEMINI LLM ----------
 class GeminiLLM(LLM):
-    ...
-    # ← aquí ya tienes __init__, metadata, chat, achat, complete, acomplete
-    # ⬇️  pega estos métodos faltantes ⬇️
+    """Wrapper mínimo para Gemini 1.5 que satisface la interfaz de Llama-Index."""
 
+    model_name: str = Field(default="models/gemini-1.5-flash-latest")
+    temperature: float = Field(default=0.0)
+
+    # -- inicialización -----------------------------------------------------
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        api_key = os.getenv("GEMINI_API_KEY")
+        if not api_key:
+            raise ValueError("GEMINI_API_KEY no configurada en variables de entorno")
+        genai.configure(api_key=api_key)
+
+        self._model = genai.GenerativeModel(
+            model_name=self.model_name,
+            generation_config=genai.types.GenerationConfig(
+                temperature=float(self.temperature)
+            ),
+        )
+
+        # callback manager defensivo
+        if self.callback_manager is None:
+            from llama_index.core.callbacks.base import CallbackManager
+            self.callback_manager = CallbackManager([])
+        if not self.callback_manager.handlers:
+            self.callback_manager.add_handler(LlamaDebugHandler())
+
+    # -- metadatos ----------------------------------------------------------
+    @property
+    def metadata(self) -> LLMMetadata:  # type: ignore[override]
+        return LLMMetadata(
+            context_window=1_048_576,
+            num_output=8192,
+            is_chat_model=True,
+            is_function_calling_model=True,
+            model_name=self.model_name,
+        )
+
+    # ----------------------------------------------------------------------
+    # 1️⃣  CHAT SINCRONO
+    # ----------------------------------------------------------------------
+    def chat(self, messages: list[ChatMessage], **kwargs) -> ChatMessage:  # type: ignore[override]
+        history = [
+            {"role": "user" if m.role == "user" else "model", "parts": [{"text": str(m.content)}]}
+            for m in messages[:-1]
+        ]
+        session = self._model.start_chat(history=history)
+        reply = session.send_message(str(messages[-1].content))
+        return ChatMessage(role="assistant", content=reply.text)
+
+    # 1-bis  CHAT ASINCRONO
+    async def achat(self, messages: list[ChatMessage], **kwargs) -> ChatMessage:  # type: ignore[override]
+        return await asyncio.to_thread(self.chat, messages, **kwargs)
+
+    # ----------------------------------------------------------------------
+    # 2️⃣  COMPLETE SINCRONO (prompt plano)
+    # ----------------------------------------------------------------------
+    def complete(self, prompt: str, formatted: bool = False, **kwargs) -> CompletionResponse:  # type: ignore[override]
+        resp = self._model.generate_content(prompt)
+        return CompletionResponse(text=resp.text)
+
+    # 2-bis COMPLETE ASINCRONO
+    async def acomplete(self, prompt: str, formatted: bool = False, **kwargs) -> CompletionResponse:  # type: ignore[override]
+        return await asyncio.to_thread(self.complete, prompt, formatted=formatted, **kwargs)
+
+    # ----------------------------------------------------------------------
+    # 3️⃣  STREAMING DE COMPLETIONS
+    # ----------------------------------------------------------------------
     def stream_complete(self, prompt: str, formatted: bool = False, **kwargs):
-        """Devuelve un generador incremental de CompletionResponse."""
-        stream = self._model.generate_content(str(prompt), stream=True)
+        stream = self._model.generate_content(prompt, stream=True)
 
-        def gen():
+        def generator():
+            from llama_index.core.llms import CompletionResponse
             acc = ""
-            from llama_index.core.llms import CompletionResponse  # import local
             for chunk in stream:
-                delta = getattr(chunk, "text", "") or (
-                    chunk.parts[0].text if getattr(chunk, "parts", None) else ""
-                )
+                delta = getattr(chunk, "text", "") or (chunk.parts[0].text if chunk.parts else "")
                 if delta:
                     acc += delta
                     yield CompletionResponse(text=acc, delta=delta)
 
-        return gen()
+        return generator()
 
     async def astream_complete(self, prompt: str, formatted: bool = False, **kwargs):
-        # ejecuta la versión síncrona en un hilo
         sync_gen = await asyncio.to_thread(self.stream_complete, prompt, formatted=formatted, **kwargs)
 
-        async def async_gen():
+        async def agen():
             for item in sync_gen:
                 yield item
 
-        return async_gen()
+        return agen()
 
+    # ----------------------------------------------------------------------
+    # 4️⃣  STREAMING DE CHAT
+    # ----------------------------------------------------------------------
     def stream_chat(self, messages: list[ChatMessage], **kwargs):
-        hist = [
+        history = [
             {"role": "user" if m.role == "user" else "model", "parts": [{"text": str(m.content)}]}
             for m in messages[:-1]
         ]
-        last = str(messages[-1].content)
-        session = self._model.start_chat(history=hist)
-        stream = session.send_message(last, stream=True)
+        session = self._model.start_chat(history=history)
+        stream = session.send_message(str(messages[-1].content), stream=True)
 
-        def gen():
+        def generator():
             acc = ""
             for chunk in stream:
-                delta = getattr(chunk, "text", "") or (
-                    chunk.parts[0].text if getattr(chunk, "parts", None) else ""
-                )
+                delta = getattr(chunk, "text", "") or (chunk.parts[0].text if chunk.parts else "")
                 if delta:
                     acc += delta
-                    yield ChatMessage(role="assistant", content=acc, additional_kwargs={"delta": delta})
+                    yield ChatMessage(
+                        role="assistant",
+                        content=acc,
+                        additional_kwargs={"delta": delta},
+                    )
 
-        return gen()
+        return generator()
 
     async def astream_chat(self, messages: list[ChatMessage], **kwargs):
         sync_gen = await asyncio.to_thread(self.stream_chat, messages, **kwargs)
 
-        async def async_gen():
+        async def agen():
             for item in sync_gen:
                 yield item
 
-        return async_gen()
+        return agen()
+