updated

Signed-off-by: Francisco Javier Arceo <[email protected]>

Author: franciscojavierarceo

llama_stack/providers/inline/tool_runtime/rag/memory.py

@@ -260,44 +260,21 @@ async def list_runtime_tools(
         )
 
     async def invoke_tool(self, tool_name: str, kwargs: dict[str, Any]) -> ToolInvocationResult:
-        if tool_name == "insert_into_memory":
-            documents = kwargs.get("documents", [])
-            vector_db_id = kwargs.get("vector_db_id")
-            chunk_size_in_tokens = kwargs.get("chunk_size_in_tokens", 512)
-
-            if not vector_db_id:
-                raise ValueError("vector_db_id is required for insert_into_memory")
-
-            await self.insert(
-                documents=documents,
-                vector_db_id=vector_db_id,
-                chunk_size_in_tokens=chunk_size_in_tokens,
-            )
-
-            return ToolInvocationResult(
-                content=[],
-                metadata={"inserted_documents": len(documents)},
-            )
-
-        elif tool_name == "knowledge_search":
-            vector_db_ids = kwargs.get("vector_db_ids", [])
-            query_config = kwargs.get("query_config")
-            if query_config:
-                query_config = TypeAdapter(RAGQueryConfig).validate_python(query_config)
-            else:
-                query_config = RAGQueryConfig()
-
-            query = kwargs["query"]
-            result = await self.query(
-                content=query,
-                vector_db_ids=vector_db_ids,
-                query_config=query_config,
-            )
+        vector_db_ids = kwargs.get("vector_db_ids", [])
+        query_config = kwargs.get("query_config")
+        if query_config:
+            query_config = TypeAdapter(RAGQueryConfig).validate_python(query_config)
+        else:
+            query_config = RAGQueryConfig()
 
-            return ToolInvocationResult(
-                content=result.content or [],
-                metadata=result.metadata,
-            )
+        query = kwargs["query"]
+        result = await self.query(
+            content=query,
+            vector_db_ids=vector_db_ids,
+            query_config=query_config,
+        )
 
-        else:
-            raise ValueError(f"Unknown tool: {tool_name}")
+        return ToolInvocationResult(
+            content=result.content or [],
+            metadata=result.metadata,
+        )