Small bugfixes on prompts, return types

js/sdk/src/r2rClient.ts

@@ -1156,7 +1156,7 @@ export class r2rClient {
     rag_generation_config?: GenerationConfig | Record<string, any>,
     task_prompt_override?: string,
     include_title_if_available?: boolean,
-  ): Promise<any> {
+  ): Promise<any | AsyncGenerator<string, void, unknown>> {
     this._ensureAuthenticated();
 
     const json_data: Record<string, any> = {
@@ -1180,38 +1180,19 @@ export class r2rClient {
   }
 
   // TODO: can we remove this and pull this into rag?
-  @featureGenerator("streamingRag")
-  private async *streamRag(
+  @feature("streamingRag")
+  private async streamRag(
     rag_data: Record<string, any>,
-  ): AsyncGenerator<string, void, unknown> {
+  ): Promise<ReadableStream<Uint8Array>> {
     this._ensureAuthenticated();
 
-    const response = await this._makeRequest<Response>("POST", "rag", {
+    return this._makeRequest<ReadableStream<Uint8Array>>("POST", "rag", {
       data: rag_data,
       headers: {
         "Content-Type": "application/json",
       },
       responseType: "stream",
     });
-
-    if (!response.body) {
-      throw new Error("Response body is null");
-    }
-
-    const reader = response.body.getReader();
-    const decoder = new TextDecoder();
-
-    try {
-      while (true) {
-        const { done, value } = await reader.read();
-        if (done) {
-          break;
-        }
-        yield decoder.decode(value);
-      }
-    } finally {
-      reader.releaseLock();
-    }
   }
 
   /**
@@ -1232,7 +1213,7 @@ export class r2rClient {
     rag_generation_config?: GenerationConfig | Record<string, any>,
     task_prompt_override?: string,
     include_title_if_available?: boolean,
-  ): Promise<any> {
+  ): Promise<any | AsyncGenerator<string, void, unknown>> {
     this._ensureAuthenticated();
 
     const json_data: Record<string, any> = {
@@ -1256,38 +1237,19 @@ export class r2rClient {
   }
 
   // TODO: can we remove this and pull this into agent?
-  @featureGenerator("streamingAgent")
-  private async *streamAgent(
+  @feature("streamingAgent")
+  private async streamAgent(
     agent_data: Record<string, any>,
-  ): AsyncGenerator<string, void, unknown> {
+  ): Promise<ReadableStream<Uint8Array>> {
     this._ensureAuthenticated();
 
-    const response = await this._makeRequest<Response>("POST", "agent", {
+    return this._makeRequest<ReadableStream<Uint8Array>>("POST", "agent", {
       data: agent_data,
       headers: {
         "Content-Type": "application/json",
       },
       responseType: "stream",
     });
-
-    if (!response.body) {
-      throw new Error("Response body is null");
-    }
-
-    const reader = response.body.getReader();
-    const decoder = new TextDecoder();
-
-    try {
-      while (true) {
-        const { done, value } = await reader.read();
-        if (done) {
-          break;
-        }
-        yield decoder.decode(value);
-      }
-    } finally {
-      reader.releaseLock();
-    }
   }
 }
 

py/cli/commands/server.py

@@ -217,7 +217,7 @@ def serve(
             ).replace(":", "")
 
     if docker:
-        
+
         run_docker_serve(
             client,
             host,

py/core/base/abstractions/__init__.py

@@ -31,8 +31,8 @@
 from .restructure import KGEnrichmentSettings
 from .search import (
     AggregateSearchResult,
-    KGLocalSearchResult,
     KGGlobalSearchResult,
+    KGLocalSearchResult,
     KGSearchResult,
     KGSearchSettings,
     VectorSearchResult,

py/core/base/abstractions/search.py

@@ -54,8 +54,10 @@ class Config:
             },
         }
 
+
 class KGLocalSearchResult(BaseModel):
     """Result of a local knowledge graph search operation."""
+
     query: str
     entities: dict[str, Any]
     relationships: dict[str, Any]
@@ -70,6 +72,7 @@ def __repr__(self) -> str:
 
 class KGGlobalSearchResult(BaseModel):
     """Result of a global knowledge graph search operation."""
+
     query: str
     search_result: list[str]
 
@@ -80,26 +83,29 @@ def __repr__(self) -> str:
         return self.__str__()
 
     def dict(self) -> dict:
-        return {
-            "query": self.query,
-            "search_result": self.search_result
-        }
+        return {"query": self.query, "search_result": self.search_result}
+
 
 class KGSearchResult(BaseModel):
     """Result of a knowledge graph search operation."""
+
     local_result: Optional[KGLocalSearchResult] = None
     global_result: Optional[KGGlobalSearchResult] = None
-    
+
     def __str__(self) -> str:
         return f"KGSearchResult(local_result={self.local_result}, global_result={self.global_result})"
 
     def __repr__(self) -> str:
         return self.__str__()
-    
+
     def dict(self) -> dict:
         return {
-            "local_result": self.local_result.dict() if self.local_result else None,
-            "global_result": self.global_result.dict() if self.global_result else None
+            "local_result": (
+                self.local_result.dict() if self.local_result else None
+            ),
+            "global_result": (
+                self.global_result.dict() if self.global_result else None
+            ),
         }
 
 

py/core/main/engine.py

@@ -81,10 +81,6 @@ def __init__(
     async def aingest_documents(self, *args, **kwargs):
         return await self.ingestion_service.ingest_documents(*args, **kwargs)
 
-    @syncable
-    async def aupdate_documents(self, *args, **kwargs):
-        return await self.ingestion_service.update_documents(*args, **kwargs)
-
     @syncable
     async def aingest_files(self, *args, **kwargs):
         return await self.ingestion_service.ingest_files(*args, **kwargs)
@@ -109,10 +105,6 @@ async def arag(self, *args, **kwargs):
     async def arag_agent(self, *args, **kwargs):
         return await self.retrieval_service.agent(*args, **kwargs)
 
-    @syncable
-    async def aevaluate(self, *args, **kwargs):
-        return await self.retrieval_service.evaluate(*args, **kwargs)
-
     @syncable
     async def aupdate_prompt(self, *args, **kwargs):
         return await self.management_service.update_prompt(*args, **kwargs)

py/core/main/services/ingestion_service.py

@@ -321,7 +321,7 @@ async def ingest_documents(
             logger.error("All provided documents already exist.")
             raise R2RException(
                 status_code=409,
-                message="All provided documents already exist. Use the `update_documents` endpoint instead to update these documents.",
+                message="All provided documents already exist. Use the `update_files` endpoint instead to update these documents.",
             )
 
         # Insert pending document info

py/core/main/services/management_service.py

@@ -321,6 +321,18 @@ async def document_chunks(
             document_id, offset=offset, limit=limit
         )
 
+    @telemetry_event("UpdatePrompt")
+    async def update_prompt(
+        self,
+        name: str,
+        template: Optional[str] = None,
+        input_types: Optional[dict[str, str]] = None,
+    ):
+        if input_types is None:
+            input_types = {}
+        self.providers.prompt.update_prompt(name, template, input_types)
+        return {"message": f"Prompt '{name}' updated successfully."}
+
     @telemetry_event("UsersOverview")
     async def users_overview(
         self,

py/core/pipes/retrieval/kg_search_search_pipe.py

@@ -12,8 +12,8 @@
     KGSearchSettings,
     PipeType,
     PromptProvider,
+    R2RException,
     RunLoggingSingleton,
-    R2RException
 )
 from core.base.abstractions.search import (
     KGGlobalSearchResult,
@@ -25,6 +25,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 class KGSearchSearchPipe(GeneratorPipe):
     """
     Embeds and stores documents using a specified embedding model and database.
@@ -132,11 +133,18 @@ async def local_search(
                 )
                 all_search_results.append(search_result)
 
-
-            if len(all_search_results[0])==0:
-                raise R2RException("No search results found. Please make sure you have run the KG enrichment step before running the search: r2r enrich-graph", 400)
+            if len(all_search_results[0]) == 0:
+                raise R2RException(
+                    "No search results found. Please make sure you have run the KG enrichment step before running the search: r2r enrich-graph",
+                    400,
+                )
 
-            yield KGLocalSearchResult(query=message, entities=all_search_results[0], relationships=all_search_results[1], communities=all_search_results[2])
+            yield KGLocalSearchResult(
+                query=message,
+                entities=all_search_results[0],
+                relationships=all_search_results[1],
+                communities=all_search_results[2],
+            )
 
     async def global_search(
         self,

py/core/providers/database/vecs/collection.py

@@ -800,11 +800,8 @@ def parse_condition(key, value):
             else:
                 # Handle JSON-based filters
                 json_col = self.table.c.metadata
-                if not key.startswith("metadata."):
-                    raise FilterError(
-                        "metadata key must start with 'metadata.'"
-                    )
-                key = key.split("metadata.")[1]
+                if key.startswith("metadata."):
+                    key.split("metadata.")[1]
                 if isinstance(value, dict):
                     if len(value) > 1:
                         raise FilterError("only one operator permitted")

py/core/providers/prompts/defaults/rag_agent.yaml

@@ -6,7 +6,7 @@ rag_agent:
     When asked a question, perform a search to find relevant information and provide a response.
 
 
-    The response should contain line-item attributions to relevent search results, and be as informative if possible.
+    The response should contain line-item attributions to relevant search results, and be as informative if possible.
 
     If no relevant results are found, then state that no results were found.
 

py/sdk/models.py

@@ -187,6 +187,7 @@ class Config:
             },
         }
 
+
 class KGLocalSearchResult(BaseModel):
     query: str
     entities: list[dict[str, Any]]
@@ -195,15 +196,16 @@ class KGLocalSearchResult(BaseModel):
 
     def __str__(self) -> str:
         return f"KGLocalSearchResult(query={self.query}, entities={self.entities}, relationships={self.relationships}, communities={self.communities})"
-    
+
     def dict(self) -> dict:
         return {
             "query": self.query,
             "entities": self.entities,
             "relationships": self.relationships,
-            "communities": self.communities
+            "communities": self.communities,
         }
 
+
 class KGGlobalSearchResult(BaseModel):
     query: str
     search_result: list[str]
@@ -213,12 +215,9 @@ def __str__(self) -> str:
 
     def __repr__(self) -> str:
         return self.__str__()
-    
+
     def dict(self) -> dict:
-        return {
-            "query": self.query,
-            "search_result": self.search_result
-        }
+        return {"query": self.query, "search_result": self.search_result}
 
 
 class KGSearchResult(BaseModel):
@@ -230,11 +229,15 @@ def __str__(self) -> str:
 
     def __repr__(self) -> str:
         return self.__str__()
-    
+
     def dict(self) -> dict:
         return {
-            "local_result": self.local_result.dict() if self.local_result else None,
-            "global_result": self.global_result.dict() if self.global_result else None
+            "local_result": (
+                self.local_result.dict() if self.local_result else None
+            ),
+            "global_result": (
+                self.global_result.dict() if self.global_result else None
+            ),
         }
 
     class Config:
@@ -245,7 +248,7 @@ class Config:
                     "entities": {
                         "Paris": {
                             "name": "Paris",
-                            "description": "Paris is the capital of France."
+                            "description": "Paris is the capital of France.",
                         }
                     },
                     "relationships": {},
@@ -255,11 +258,12 @@ class Config:
                     "query": "What is the capital of France?",
                     "search_result": [
                         "Paris is the capital and most populous city of France."
-                    ]
-                }
+                    ],
+                },
             }
         }
 
+
 class R2RException(Exception):
     def __init__(
         self, message: str, status_code: int, detail: Optional[Any] = None

py/tests/regression/expected_outputs/test_document_management.json

@@ -203,7 +203,7 @@
     }
   },
   "reingest_sample_file": {
-    "results": "All provided documents already exist. Use the `update_documents` endpoint instead to update these documents."
+    "results": "All provided documents already exist. Use the `update_files` endpoint instead to update these documents."
   },
   "documents_overview": {
     "results": [