Merge branch 'main' into RAG_secret

kaito-project · Feb 25, 2025 · 9229f8a · 9229f8a
2 parents 70e3047 + 45e6882
commit 9229f8a
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/presets/ragengine/inference/inference.py b/presets/ragengine/inference/inference.py
@@ -89,8 +89,8 @@ async def _async_openai_complete(self, prompt: str, **kwargs: Any) -> Completion
         return await OpenAI(api_key=LLM_ACCESS_SECRET, **kwargs).acomplete(prompt)
 
     async def _async_huggingface_remote_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
-        return await self._async_post_request({"messages": [{"role": "user", "content": prompt}]}, headers={"Authorization": f"Bearer {LLM_ACCESS_SECRET}"})
-
+        data = {"prompt": prompt, **kwargs}
+        return await self._async_post_request(data, headers={"Authorization": f"Bearer {LLM_ACCESS_SECRET}", "Content-Type": "application/json"})
     async def _async_custom_api_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         model_name, model_max_len = self._get_default_model_info()
         if kwargs.get("model"):

diff --git a/presets/ragengine/models.py b/presets/ragengine/models.py
@@ -28,7 +28,7 @@ class IndexRequest(BaseModel):
 class QueryRequest(BaseModel):
     index_name: str
     query: str
-    top_k: int = 10
+    top_k: int = 5
     # Accept a dictionary for our LLM parameters
     llm_params: Optional[Dict[str, Any]] = Field(
         default_factory=dict,
@@ -70,4 +70,4 @@ class QueryResponse(BaseModel):
 
 class HealthStatus(BaseModel):
     status: str
-    detail: Optional[str] = None 
+    detail: Optional[str] = None 
diff --git a/presets/ragengine/vector_store/base.py b/presets/ragengine/vector_store/base.py
@@ -191,7 +191,7 @@ async def add_document_to_index(self, index_name: str, document: Document, doc_i
         """Common logic for adding a single document."""
         if index_name not in self.index_map:
             raise ValueError(f"No such index: '{index_name}' exists.")
-        llama_doc = LlamaDocument(id_=doc_id, text=document.text, metadata=document.metadata)
+        llama_doc = LlamaDocument(id_=doc_id, text=document.text, metadata=document.metadata, excluded_llm_metadata_keys=[key for key in document.metadata.keys()])
 
         if self.use_rwlock:
             async with self.rwlock.writer_lock: