jj

my_tools.py

@@ -41,6 +41,11 @@ class GeminiLLM(LLM):
             model_name=self.model_name,
             generation_config=self._gen_cfg
         )
+        # Inicializar callback_manager si no se pasa en kwargs
+        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())
 
@@ -54,7 +59,7 @@ class GeminiLLM(LLM):
             model_name=self.model_name,
         )
 
-    def chat(self, messages, **kwargs):
+    def chat(self, messages: list[ChatMessage], **kwargs): # Añadido tipo para messages
         hist = []
         for m in messages[:-1]:
             role = "user" if m.role == "user" else "model"
@@ -65,12 +70,12 @@ class GeminiLLM(LLM):
             resp = session.send_message(last)
             return ChatMessage(role="assistant", content=resp.text)
         except Exception as e:
-            return ChatMessage(role="assistant", content=f"Error Gemini: {e}")
+            return ChatMessage(role="assistant", content=f"Error Gemini chat: {e}")
 
-    async def achat(self, messages, **kwargs):
+    async def achat(self, messages: list[ChatMessage], **kwargs): # Añadido tipo para messages
         return await asyncio.to_thread(self.chat, messages, **kwargs)
 
-    def stream_complete(self, prompt, formatted=False, **kwargs):
+    def stream_complete(self, prompt: str, formatted=False, **kwargs): # Añadido tipo para prompt
         stream = self._model.generate_content(str(prompt), stream=True)
         def gen():
             acc = ""
@@ -83,10 +88,17 @@ class GeminiLLM(LLM):
                     yield CompletionResponse(text=acc, delta=delta)
         return gen()
 
-    async def astream_complete(self, prompt, formatted=False, **kwargs):
-        return await asyncio.to_thread(self.stream_complete, prompt, formatted=formatted, **kwargs)
+    async def astream_complete(self, prompt: str, formatted=False, **kwargs): # Añadido tipo para prompt
+        # Correctamente, esto debería devolver un generador asíncrono.
+        # Envolver el generador síncrono es un workaround común.
+        sync_gen = await asyncio.to_thread(self.stream_complete, prompt, formatted=formatted, **kwargs)
+        async def async_gen_wrapper():
+            for item in sync_gen:
+                yield item
+        return async_gen_wrapper()
 
-    def stream_chat(self, messages, **kwargs):
+
+    def stream_chat(self, messages: list[ChatMessage], **kwargs): # Añadido tipo para messages
         hist = []
         for m in messages[:-1]:
             role = "user" if m.role == "user" else "model"
@@ -105,16 +117,25 @@ class GeminiLLM(LLM):
                     yield ChatMessage(role="assistant", content=acc, additional_kwargs={"delta": delta})
         return gen()
 
-    def complete(self, prompt, formatted=False, **kwargs):
+    # --- MÉTODO FALTANTE AÑADIDO AQUÍ ---
+    async def astream_chat(self, messages: list[ChatMessage], **kwargs): # Añadido tipo para messages
+        # Similar a astream_complete, envolvemos el generador síncrono
+        sync_gen = await asyncio.to_thread(self.stream_chat, messages, **kwargs)
+        async def async_gen_wrapper():
+            for item in sync_gen:
+                yield item
+        return async_gen_wrapper()
+    # --- FIN DEL MÉTODO AÑADIDO ---
+
+    def complete(self, prompt: str, formatted=False, **kwargs): # Añadido tipo para prompt
         try:
             resp = self._model.generate_content(str(prompt))
             return CompletionResponse(text=resp.text)
         except Exception as e:
-            return CompletionResponse(text=f"Error complete: {e}")
+            return CompletionResponse(text=f"Error Gemini complete: {e}")
 
-    async def acomplete(self, prompt, formatted=False, **kwargs):
+    async def acomplete(self, prompt: str, formatted=False, **kwargs): # Añadido tipo para prompt
         return await asyncio.to_thread(self.complete, prompt, formatted=formatted, **kwargs)
-
 # -------------------------------------------------------------------
 # 2) Herramientas
 # -------------------------------------------------------------------