fix: response

services/model_handler.py

@@ -9,6 +9,14 @@ from tenacity import retry, stop_after_attempt, wait_exponential
 
 MODEL_PATH = "google/flan-t5-small"
 
+# Simple Response class to wrap the model output
+class Response:
+    def __init__(self, content):
+        self.content = content
+    
+    def __str__(self):
+        return self.content
+
 # Personnalized class for local models
 class LocalHuggingFaceModel(Model):
     def __init__(self, model, tokenizer, max_length=512):
@@ -62,31 +70,33 @@ class LocalHuggingFaceModel(Model):
         """Parse the provider response delta for streaming"""
         return delta
         
-    async def aresponse(self, prompt=None, **kwargs) -> str:
+    async def aresponse(self, prompt=None, **kwargs):
         """Async response method - required abstract method"""
         if prompt is None:
             prompt = kwargs.get('input', '')
-        return await self.ainvoke(prompt=prompt, **kwargs)
+        content = await self.ainvoke(prompt=prompt, **kwargs)
+        return Response(content)
     
     async def aresponse_stream(self, prompt=None, **kwargs):
         """Async streaming response method - required abstract method"""
         if prompt is None:
             prompt = kwargs.get('input', '')
         async for chunk in self.ainvoke_stream(prompt=prompt, **kwargs):
-            yield chunk
+            yield Response(chunk)
     
-    def response(self, prompt=None, **kwargs) -> str:
+    def response(self, prompt=None, **kwargs):
         """Synchronous response method - required abstract method"""
         if prompt is None:
             prompt = kwargs.get('input', '')
-        return self.invoke(prompt=prompt, **kwargs)
+        content = self.invoke(prompt=prompt, **kwargs)
+        return Response(content)
     
     def response_stream(self, prompt=None, **kwargs):
         """Synchronous streaming response method - required abstract method"""
         if prompt is None:
             prompt = kwargs.get('input', '')
         for chunk in self.invoke_stream(prompt=prompt, **kwargs):
-            yield chunk
+            yield Response(chunk)
         
     def generate(self, prompt: str, **kwargs):
         try:
@@ -243,7 +253,6 @@ Output:"""
             
             # Get English translation
             translation = self.translator.run(prompt=translation_prompt, stream=False)
-            logging.info(f"Translation: {translation}")
             if not translation:
                 logging.error("Translation failed")
                 return "Error: Unable to translate the query"