minor fixes

meta-llama · Dec 16, 2024 · 72c2ade · 72c2ade
1 parent 095ae36
commit 72c2ade
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 0 deletions.
diff --git a/llama_stack/apis/tools/tools.py b/llama_stack/apis/tools/tools.py
@@ -61,6 +61,7 @@ async def register_tool(
         parameters: List[ToolParameter],
         returns: ToolReturn,
         provider_id: Optional[str] = None,
+        provider_resource_id: Optional[str] = None,
         provider_metadata: Optional[Dict[str, Any]] = None,
         tool_prompt_format: Optional[ToolPromptFormat] = None,
     ) -> Tool:

diff --git a/llama_stack/distribution/datatypes.py b/llama_stack/distribution/datatypes.py
@@ -21,6 +21,7 @@
 from llama_stack.apis.memory import Memory
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.scoring import Scoring
+from llama_stack.apis.tools import Tool, ToolRuntime
 from llama_stack.providers.utils.kvstore.config import KVStoreConfig
 
 LLAMA_STACK_BUILD_CONFIG_VERSION = "2"
@@ -37,6 +38,7 @@
     Dataset,
     ScoringFn,
     EvalTask,
+    Tool,
 ]
 
 
@@ -48,6 +50,7 @@
         Dataset,
         ScoringFn,
         EvalTask,
+        Tool,
     ],
     Field(discriminator="type"),
 ]
@@ -59,6 +62,7 @@
     DatasetIO,
     Scoring,
     Eval,
+    ToolRuntime,
 ]
 
 

diff --git a/llama_stack/distribution/routers/routing_tables.py b/llama_stack/distribution/routers/routing_tables.py
@@ -485,6 +485,7 @@ async def register_tool(
         parameters: List[ToolParameter],
         returns: ToolReturn,
         provider_id: Optional[str] = None,
+        provider_resource_id: Optional[str] = None,
         provider_metadata: Optional[Dict[str, Any]] = None,
         tool_prompt_format: Optional[ToolPromptFormat] = None,
     ) -> None:
@@ -499,13 +500,17 @@ async def register_tool(
                 raise ValueError(
                     "No provider specified and multiple providers available. Please specify a provider_id."
                 )
+        if provider_resource_id is None:
+            provider_resource_id = tool_id
+
         tool = Tool(
             identifier=tool_id,
             name=name,
             description=description,
             parameters=parameters,
             returns=returns,
             provider_id=provider_id,
+            provider_resource_id=provider_resource_id,
             provider_metadata=provider_metadata,
             tool_prompt_format=tool_prompt_format,
         )