Merge pull request #780 from apecloud/refine_config

Refine config

Author: earayu

aperag/apps.py

@@ -33,6 +33,9 @@ def ready(self):
         from django.conf import settings
         settings.AUTH_USER_MODEL = 'aperag.User'
 
+        from aperag.llm.litellm_track import register_llm_track
+        register_llm_track()
+
 
 def get_ip_config():
     from django.db import transaction

aperag/embed/embedding_service.py

@@ -83,5 +83,6 @@ def _embed_batch(self, batch: Sequence[str]) -> List[List[float]]:
             api_base=self.api_base,
             api_key=self.api_key,
             input=list(batch),
+            max_retries=0,
         )
         return [item["embedding"] for item in response["data"]]

aperag/llm/completion_service.py

@@ -33,11 +33,12 @@ def provide_default_token():
         return False
 
     async def _agenerate_stream(self, history, prompt, memory=False):
-        response = litellm.completion(
+        response = await litellm.acompletion(
             **self.kwargs,
             messages=history + [{"role": "user", "content": prompt}] if memory else [{"role": "user", "content": prompt}],
             stream=True,
             timeout=httpx.Timeout(None, connect=3),
+            max_retries=0,
         )
         async for chunk in response:
             if not chunk.choices:
@@ -53,6 +54,7 @@ def _generate_stream(self, history, prompt, memory=False):
             messages=history + [{"role": "user", "content": prompt}] if memory else [{"role": "user", "content": prompt}],
             stream=True,
             timeout=httpx.Timeout(None, connect=3),
+            max_retries=0,
         )
         for chunk in response:
             if not chunk.choices:
@@ -67,11 +69,12 @@ async def agenerate_stream(self, history, prompt, memory=False):
             yield tokens
 
     async def agenerate_by_tools(self, prompt,tools):
-        response = litellm.completion(
+        response = await litellm.acompletion(
             **self.kwargs,
             messages=[{"role": "user", "content": prompt}],
             tools=tools,
             tool_choice="auto",
+            max_retries=0,
         )
         tool_calls = response.choices[0].message.tool_calls
         return tool_calls, response.choices[0].message.content

aperag/llm/litellm_track.py

@@ -1,11 +1,13 @@
 import logging
 import litellm
 from datetime import datetime
+from litellm.integrations.custom_logger import CustomLogger
+
 
 logger = logging.getLogger(__name__)
 
 
-class MyCustomHandler:
+class MyCustomHandler(CustomLogger):
     def __init__(self):
         self.prefix = "LiteLLM"
 
@@ -17,30 +19,13 @@ def _format_time(self, dt: datetime) -> str:
     def log_pre_api_call(self, model, messages, kwargs):
         pass
 
+    # async func won't trigger this callback, it's a bug. https://github.com/BerriAI/litellm/issues/8842
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
-        try:
-            latency = (end_time - start_time).total_seconds() if start_time and end_time else 0.0
-            cost = kwargs.get("response_cost", 0)
-
-            logger.info(f"{self.prefix} Status: Success")
-            logger.info(f"{self.prefix} Model: {kwargs.get('model')}")
-            logger.info(f"{self.prefix} Provider: {kwargs.get('custom_llm_provider', 'unknown')}")
-            logger.info(f"{self.prefix} Start Time: {self._format_time(start_time)}")
-            logger.info(f"{self.prefix} End Time: {self._format_time(end_time)}")
-            logger.info(f"{self.prefix} Latency: {latency:.2f} seconds")
-            logger.info(f"{self.prefix} Cost: ${cost:.6f}")
-        except Exception as e:
-            logger.error(f"{self.prefix} Error in log_success_event", exc_info=e)
+        pass
 
     def log_failure_event(self, kwargs, response_obj, start_time, end_time):
         try:
             error_str = str(response_obj) if response_obj else "Unknown error"
-
-            logger.error(f"{self.prefix} Status: Failed")
-            logger.error(f"{self.prefix} Model: {kwargs.get('model')}")
-            logger.error(f"{self.prefix} Provider: {kwargs.get('custom_llm_provider', 'unknown')}")
-            logger.error(f"{self.prefix} Start Time: {self._format_time(start_time)}")
-            logger.error(f"{self.prefix} End Time: {self._format_time(end_time)}")
             logger.error(f"{self.prefix} Error: {error_str}", exc_info=True)
         except Exception as e:
             logger.error(f"{self.prefix} Error in log_failure_event", exc_info=e)
@@ -52,7 +37,19 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti
         self.log_failure_event(kwargs, response_obj, start_time, end_time)
 
     def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
-        pass
+        try:
+            latency = (end_time - start_time).total_seconds() if start_time and end_time else 0.0
+            cost = kwargs.get("response_cost", 0) or 0
+
+            logger.info(f"{self.prefix} Model: {kwargs.get('model')}")
+            logger.info(f"{self.prefix} Provider: {kwargs.get('custom_llm_provider', 'unknown')}")
+            logger.info(f"{self.prefix} base_url: {kwargs.get('api_base', 'unknown')}")
+            logger.info(f"{self.prefix} Start Time: {self._format_time(start_time)}")
+            logger.info(f"{self.prefix} End Time: {self._format_time(end_time)}")
+            logger.info(f"{self.prefix} Latency: {latency:.2f} seconds")
+            logger.info(f"{self.prefix} Cost: ${cost:.6f}")
+        except Exception as e:
+            logger.error(f"{self.prefix} Error in log_success_event", exc_info=e)
 
 
 def register_llm_track():

aperag/rank/reranker.py

@@ -29,7 +29,8 @@ async def rank(self, query: str, results: List[DocumentWithScore]):
             top_n=len(documents),
             api_key=self.api_key,
             api_base=self.api_base,
-            return_documents=False
+            return_documents=False,
+            max_retries=0,
         )
         indices = [item["index"] for item in resp["results"]]
         return [results[i] for i in indices]

tests/unit_test/embed/test_embedding_service.py

@@ -1,6 +1,7 @@
 import os
 import pytest
 from aperag.embed.embedding_service import EmbeddingService
+from aperag.llm.litellm_track import register_llm_track
 
 PROVIDERS = [
     {
@@ -38,6 +39,9 @@
 ]
 
 
+register_llm_track()
+
+
 @pytest.fixture(params=PROVIDERS)
 def embedding_service(request):
     config = request.param