Upload 3 files

app/services/enrichment.py

@@ -4,8 +4,8 @@ from litellm import completion
 from app.core.config import settings
 
 class AIEnrichmentService:
-    def __init__(self, model: str = settings.DEFAULT_LLM_MODEL):
-        self.model = model
+    def __init__(self, model: Optional[str] = None):
+        self.model = model or settings.get_effective_llm()
         self._local_embedding_model = None
 
     def _get_local_embedding_model(self):
@@ -42,6 +42,7 @@ class AIEnrichmentService:
             response = await completion(
                 model=self.model,
                 messages=[{"role": "user", "content": prompt}],
+                api_key=settings.OPENAI_API_KEY or settings.HUGGINGFACE_API_TOKEN,
                 response_format={"type": "json_object"}
             )
             import json
@@ -52,17 +53,20 @@ class AIEnrichmentService:
             return {"entities": [], "tags": [], "sensitivity": "unknown", "pii_detected": False}
 
     async def generate_embedding(self, text: str) -> List[float]:
-        # 1. Try OpenAI if key is present
-        if settings.OPENAI_API_KEY:
+        effective_model = settings.get_effective_embedding()
+        
+        # 1. Try API if key is present
+        if settings.OPENAI_API_KEY or (settings.HUGGINGFACE_API_TOKEN and "huggingface" in effective_model):
             try:
                 from litellm import embedding
                 response = await embedding(
-                    model=settings.DEFAULT_EMBEDDING_MODEL,
-                    input=[text]
+                    model=effective_model,
+                    input=[text],
+                    api_key=settings.OPENAI_API_KEY or settings.HUGGINGFACE_API_TOKEN
                 )
                 return response['data'][0]['embedding']
             except Exception as e:
-                print(f"OpenAI embedding failed, trying fallback: {e}")
+                print(f"API embedding ({effective_model}) failed, trying fallback: {e}")
 
         # 2. Fallback to Local Sentence Transformers (CPU friendly)
         try:

app/services/generation.py

@@ -1,10 +1,11 @@
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
+
 from litellm import completion
 from app.core.config import settings
 
 class GenerationService:
-    def __init__(self, model: str = settings.DEFAULT_LLM_MODEL):
-        self.model = model
+    def __init__(self, model: Optional[str] = None):
+        self.model = model or settings.get_effective_llm()
 
     async def generate_answer(self, query: str, context_chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
         context_text = "\n\n".join([
@@ -35,8 +36,9 @@ class GenerationService:
             }
 
         try:
+            effective_model = self.model # Already adjusted in __init__
             response = await completion(
-                model=self.model,
+                model=effective_model,
                 messages=[{"role": "user", "content": prompt}],
                 api_key=settings.OPENAI_API_KEY or settings.HUGGINGFACE_API_TOKEN
             )

app/services/retrieval.py

@@ -15,9 +15,10 @@ class RetrievalService:
             return query
             
         try:
+            effective_model = settings.get_effective_llm()
             prompt = f"Rewrite the following user query to be more search-friendly for a RAG system. Original: {query}"
             response = await completion(
-                model=settings.DEFAULT_LLM_MODEL,
+                model=effective_model,
                 messages=[{"role": "user", "content": prompt}],
                 api_key=settings.OPENAI_API_KEY or settings.HUGGINGFACE_API_TOKEN
             )