Add back tast prompt override and include title if availible

SciPhi-AI · Aug 22, 2024 · 144bc1c · 144bc1c
1 parent 7d6c7f2
commit 144bc1c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/js/sdk/src/r2rClient.ts b/js/sdk/src/r2rClient.ts
@@ -1139,6 +1139,8 @@ export class r2rClient {
    * @param vector_search_settings Vector search settings.
    * @param kg_search_settings KG search settings.
    * @param rag_generation_config RAG generation configuration.
+   * @param task_prompt_override Task prompt override.
+   * @param include_title_if_available Include title if available.
    * @returns A promise that resolves to the response from the server.
    */
   @feature("rag")
@@ -1147,6 +1149,8 @@ export class r2rClient {
     vector_search_settings?: VectorSearchSettings | Record<string, any>,
     kg_search_settings?: KGSearchSettings | Record<string, any>,
     rag_generation_config?: GenerationConfig | Record<string, any>,
+    task_prompt_override?: string,
+    include_title_if_available?: boolean,
   ): Promise<any> {
     this._ensureAuthenticated();
 
@@ -1155,6 +1159,8 @@ export class r2rClient {
       vector_search_settings,
       kg_search_settings,
       rag_generation_config,
+      task_prompt_override,
+      include_title_if_available,
     };
 
     Object.keys(json_data).forEach(

diff --git a/py/sdk/retrieval.py b/py/sdk/retrieval.py
@@ -19,8 +19,6 @@ async def search(
             Union[dict, VectorSearchSettings]
         ] = None,
         kg_search_settings: Optional[Union[dict, KGSearchSettings]] = None,
-        *args,
-        **kwargs,
     ) -> SearchResponse:
         """
         Conduct a vector and/or KG search.
@@ -56,8 +54,8 @@ async def rag(
             Union[dict, VectorSearchSettings]
         ] = None,
         kg_search_settings: Optional[Union[dict, KGSearchSettings]] = None,
-        *args,
-        **kwargs,
+        task_prompt_override: Optional[str] = None,
+        include_title_if_available: Optional[bool] = False,
     ) -> Union[RAGResponse, AsyncGenerator[RAGResponse, None]]:
         """
         Conducts a Retrieval Augmented Generation (RAG) search with the given query.
@@ -67,6 +65,8 @@ async def rag(
             rag_generation_config (Optional[Union[dict, GenerationConfig]]): RAG generation configuration.
             vector_search_settings (Optional[Union[dict, VectorSearchSettings]]): Vector search settings.
             kg_search_settings (Optional[Union[dict, KGSearchSettings]]): KG search settings.
+            task_prompt_override (Optional[str]): Task prompt override.
+            include_title_if_available (Optional[bool]): Include the title if available.
 
         Returns:
             Union[RAGResponse, AsyncGenerator[RAGResponse, None]]: The RAG response
@@ -87,6 +87,8 @@ async def rag(
             "rag_generation_config": rag_generation_config,
             "vector_search_settings": vector_search_settings,
             "kg_search_settings": kg_search_settings,
+            "task_prompt_override": task_prompt_override,
+            "include_title_if_available": include_title_if_available,
         }
 
         if rag_generation_config and rag_generation_config.get(
@@ -107,8 +109,6 @@ async def agent(
         kg_search_settings: Optional[Union[dict, KGSearchSettings]] = None,
         task_prompt_override: Optional[str] = None,
         include_title_if_available: Optional[bool] = False,
-        *args,
-        **kwargs,
     ) -> Union[list[Message], AsyncGenerator[Message, None]]:
         """
         Performs a single turn in a conversation with a RAG agent.