minor fixes

This commit is contained in:
Dinesh Yeduguru 2024-12-16 13:01:52 -08:00
parent e5ac10f803
commit b7ad53ca93
3 changed files with 10 additions and 0 deletions

View file

@ -61,6 +61,7 @@ class Tools(Protocol):
parameters: List[ToolParameter], parameters: List[ToolParameter],
returns: ToolReturn, returns: ToolReturn,
provider_id: Optional[str] = None, provider_id: Optional[str] = None,
provider_resource_id: Optional[str] = None,
provider_metadata: Optional[Dict[str, Any]] = None, provider_metadata: Optional[Dict[str, Any]] = None,
tool_prompt_format: Optional[ToolPromptFormat] = None, tool_prompt_format: Optional[ToolPromptFormat] = None,
) -> Tool: ) -> Tool:

View file

@ -21,6 +21,7 @@ from llama_stack.apis.inference import Inference
from llama_stack.apis.memory import Memory from llama_stack.apis.memory import Memory
from llama_stack.apis.safety import Safety from llama_stack.apis.safety import Safety
from llama_stack.apis.scoring import Scoring from llama_stack.apis.scoring import Scoring
from llama_stack.apis.tools import Tool, ToolRuntime
from llama_stack.providers.utils.kvstore.config import KVStoreConfig from llama_stack.providers.utils.kvstore.config import KVStoreConfig
LLAMA_STACK_BUILD_CONFIG_VERSION = "2" LLAMA_STACK_BUILD_CONFIG_VERSION = "2"
@ -37,6 +38,7 @@ RoutableObject = Union[
Dataset, Dataset,
ScoringFn, ScoringFn,
EvalTask, EvalTask,
Tool,
] ]
@ -48,6 +50,7 @@ RoutableObjectWithProvider = Annotated[
Dataset, Dataset,
ScoringFn, ScoringFn,
EvalTask, EvalTask,
Tool,
], ],
Field(discriminator="type"), Field(discriminator="type"),
] ]
@ -59,6 +62,7 @@ RoutedProtocol = Union[
DatasetIO, DatasetIO,
Scoring, Scoring,
Eval, Eval,
ToolRuntime,
] ]

View file

@ -482,6 +482,7 @@ class ToolsRoutingTable(CommonRoutingTableImpl, Tools):
parameters: List[ToolParameter], parameters: List[ToolParameter],
returns: ToolReturn, returns: ToolReturn,
provider_id: Optional[str] = None, provider_id: Optional[str] = None,
provider_resource_id: Optional[str] = None,
provider_metadata: Optional[Dict[str, Any]] = None, provider_metadata: Optional[Dict[str, Any]] = None,
tool_prompt_format: Optional[ToolPromptFormat] = None, tool_prompt_format: Optional[ToolPromptFormat] = None,
) -> None: ) -> None:
@ -496,6 +497,9 @@ class ToolsRoutingTable(CommonRoutingTableImpl, Tools):
raise ValueError( raise ValueError(
"No provider specified and multiple providers available. Please specify a provider_id." "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( tool = Tool(
identifier=tool_id, identifier=tool_id,
name=name, name=name,
@ -503,6 +507,7 @@ class ToolsRoutingTable(CommonRoutingTableImpl, Tools):
parameters=parameters, parameters=parameters,
returns=returns, returns=returns,
provider_id=provider_id, provider_id=provider_id,
provider_resource_id=provider_resource_id,
provider_metadata=provider_metadata, provider_metadata=provider_metadata,
tool_prompt_format=tool_prompt_format, tool_prompt_format=tool_prompt_format,
) )