llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/distribution/resolver.py

@@ -1,323 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-import importlib
-import inspect
-
-from typing import Any, Dict, List, Set
-
-from llama_stack.providers.datatypes import *  # noqa: F403
-from llama_stack.distribution.datatypes import *  # noqa: F403
-
-from llama_stack.apis.agents import Agents
-from llama_stack.apis.inference import Inference
-from llama_stack.apis.inspect import Inspect
-from llama_stack.apis.memory import Memory
-from llama_stack.apis.memory_banks import MemoryBanks
-from llama_stack.apis.models import Models
-from llama_stack.apis.safety import Safety
-from llama_stack.apis.shields import Shields
-from llama_stack.apis.telemetry import Telemetry
-from llama_stack.distribution.distribution import (
-    builtin_automatically_routed_apis,
-    get_provider_registry,
-)
-from llama_stack.distribution.utils.dynamic import instantiate_class_type
-
-
-def api_protocol_map() -> Dict[Api, Any]:
-    return {
-        Api.agents: Agents,
-        Api.inference: Inference,
-        Api.inspect: Inspect,
-        Api.memory: Memory,
-        Api.memory_banks: MemoryBanks,
-        Api.models: Models,
-        Api.safety: Safety,
-        Api.shields: Shields,
-        Api.telemetry: Telemetry,
-    }
-
-
-def additional_protocols_map() -> Dict[Api, Any]:
-    return {
-        Api.inference: ModelsProtocolPrivate,
-        Api.memory: MemoryBanksProtocolPrivate,
-        Api.safety: ShieldsProtocolPrivate,
-    }
-
-
-# TODO: make all this naming far less atrocious. Provider. ProviderSpec. ProviderWithSpec. WTF!
-class ProviderWithSpec(Provider):
-    spec: ProviderSpec
-
-
-# TODO: this code is not very straightforward to follow and needs one more round of refactoring
-async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, Any]:
-    """
-    Does two things:
-    - flatmaps, sorts and resolves the providers in dependency order
-    - for each API, produces either a (local, passthrough or router) implementation
-    """
-    all_api_providers = get_provider_registry()
-
-    routing_table_apis = set(
-        x.routing_table_api for x in builtin_automatically_routed_apis()
-    )
-    router_apis = set(x.router_api for x in builtin_automatically_routed_apis())
-
-    providers_with_specs = {}
-
-    for api_str, providers in run_config.providers.items():
-        api = Api(api_str)
-        if api in routing_table_apis:
-            raise ValueError(
-                f"Provider for `{api_str}` is automatically provided and cannot be overridden"
-            )
-
-        specs = {}
-        for provider in providers:
-            if provider.provider_type not in all_api_providers[api]:
-                raise ValueError(
-                    f"Provider `{provider.provider_type}` is not available for API `{api}`"
-                )
-
-            p = all_api_providers[api][provider.provider_type]
-            p.deps__ = [a.value for a in p.api_dependencies]
-            spec = ProviderWithSpec(
-                spec=p,
-                **(provider.dict()),
-            )
-            specs[provider.provider_id] = spec
-
-        key = api_str if api not in router_apis else f"inner-{api_str}"
-        providers_with_specs[key] = specs
-
-    apis_to_serve = run_config.apis or set(
-        list(providers_with_specs.keys())
-        + [x.value for x in routing_table_apis]
-        + [x.value for x in router_apis]
-    )
-
-    for info in builtin_automatically_routed_apis():
-        if info.router_api.value not in apis_to_serve:
-            continue
-
-        available_providers = providers_with_specs[f"inner-{info.router_api.value}"]
-
-        providers_with_specs[info.routing_table_api.value] = {
-            "__builtin__": ProviderWithSpec(
-                provider_id="__routing_table__",
-                provider_type="__routing_table__",
-                config={},
-                spec=RoutingTableProviderSpec(
-                    api=info.routing_table_api,
-                    router_api=info.router_api,
-                    module="llama_stack.distribution.routers",
-                    api_dependencies=[],
-                    deps__=([f"inner-{info.router_api.value}"]),
-                ),
-            )
-        }
-
-        providers_with_specs[info.router_api.value] = {
-            "__builtin__": ProviderWithSpec(
-                provider_id="__autorouted__",
-                provider_type="__autorouted__",
-                config={},
-                spec=AutoRoutedProviderSpec(
-                    api=info.router_api,
-                    module="llama_stack.distribution.routers",
-                    routing_table_api=info.routing_table_api,
-                    api_dependencies=[info.routing_table_api],
-                    deps__=([info.routing_table_api.value]),
-                ),
-            )
-        }
-
-    sorted_providers = topological_sort(
-        {k: v.values() for k, v in providers_with_specs.items()}
-    )
-    apis = [x[1].spec.api for x in sorted_providers]
-    sorted_providers.append(
-        (
-            "inspect",
-            ProviderWithSpec(
-                provider_id="__builtin__",
-                provider_type="__builtin__",
-                config={
-                    "run_config": run_config.dict(),
-                },
-                spec=InlineProviderSpec(
-                    api=Api.inspect,
-                    provider_type="__builtin__",
-                    config_class="llama_stack.distribution.inspect.DistributionInspectConfig",
-                    module="llama_stack.distribution.inspect",
-                    api_dependencies=apis,
-                    deps__=([x.value for x in apis]),
-                ),
-            ),
-        )
-    )
-
-    print(f"Resolved {len(sorted_providers)} providers")
-    for api_str, provider in sorted_providers:
-        print(f" {api_str} => {provider.provider_id}")
-    print("")
-
-    impls = {}
-    inner_impls_by_provider_id = {f"inner-{x.value}": {} for x in router_apis}
-    for api_str, provider in sorted_providers:
-        deps = {a: impls[a] for a in provider.spec.api_dependencies}
-
-        inner_impls = {}
-        if isinstance(provider.spec, RoutingTableProviderSpec):
-            inner_impls = inner_impls_by_provider_id[
-                f"inner-{provider.spec.router_api.value}"
-            ]
-
-        impl = await instantiate_provider(
-            provider,
-            deps,
-            inner_impls,
-        )
-        # TODO: ugh slightly redesign this shady looking code
-        if "inner-" in api_str:
-            inner_impls_by_provider_id[api_str][provider.provider_id] = impl
-        else:
-            api = Api(api_str)
-            impls[api] = impl
-
-    return impls
-
-
-def topological_sort(
-    providers_with_specs: Dict[str, List[ProviderWithSpec]],
-) -> List[ProviderWithSpec]:
-    def dfs(kv, visited: Set[str], stack: List[str]):
-        api_str, providers = kv
-        visited.add(api_str)
-
-        deps = []
-        for provider in providers:
-            for dep in provider.spec.deps__:
-                deps.append(dep)
-
-        for dep in deps:
-            if dep not in visited:
-                dfs((dep, providers_with_specs[dep]), visited, stack)
-
-        stack.append(api_str)
-
-    visited = set()
-    stack = []
-
-    for api_str, providers in providers_with_specs.items():
-        if api_str not in visited:
-            dfs((api_str, providers), visited, stack)
-
-    flattened = []
-    for api_str in stack:
-        for provider in providers_with_specs[api_str]:
-            flattened.append((api_str, provider))
-    return flattened
-
-
-# returns a class implementing the protocol corresponding to the Api
-async def instantiate_provider(
-    provider: ProviderWithSpec,
-    deps: Dict[str, Any],
-    inner_impls: Dict[str, Any],
-):
-    protocols = api_protocol_map()
-    additional_protocols = additional_protocols_map()
-
-    provider_spec = provider.spec
-    module = importlib.import_module(provider_spec.module)
-
-    args = []
-    if isinstance(provider_spec, RemoteProviderSpec):
-        if provider_spec.adapter:
-            method = "get_adapter_impl"
-        else:
-            method = "get_client_impl"
-
-        config_type = instantiate_class_type(provider_spec.config_class)
-        config = config_type(**provider.config)
-        args = [config, deps]
-    elif isinstance(provider_spec, AutoRoutedProviderSpec):
-        method = "get_auto_router_impl"
-
-        config = None
-        args = [provider_spec.api, deps[provider_spec.routing_table_api], deps]
-    elif isinstance(provider_spec, RoutingTableProviderSpec):
-        method = "get_routing_table_impl"
-
-        config = None
-        args = [provider_spec.api, inner_impls, deps]
-    else:
-        method = "get_provider_impl"
-
-        config_type = instantiate_class_type(provider_spec.config_class)
-        config = config_type(**provider.config)
-        args = [config, deps]
-
-    fn = getattr(module, method)
-    impl = await fn(*args)
-    impl.__provider_id__ = provider.provider_id
-    impl.__provider_spec__ = provider_spec
-    impl.__provider_config__ = config
-
-    check_protocol_compliance(impl, protocols[provider_spec.api])
-    if (
-        not isinstance(provider_spec, AutoRoutedProviderSpec)
-        and provider_spec.api in additional_protocols
-    ):
-        additional_api = additional_protocols[provider_spec.api]
-        check_protocol_compliance(impl, additional_api)
-
-    return impl
-
-
-def check_protocol_compliance(obj: Any, protocol: Any) -> None:
-    missing_methods = []
-
-    mro = type(obj).__mro__
-    for name, value in inspect.getmembers(protocol):
-        if inspect.isfunction(value) and hasattr(value, "__webmethod__"):
-            if not hasattr(obj, name):
-                missing_methods.append((name, "missing"))
-            elif not callable(getattr(obj, name)):
-                missing_methods.append((name, "not_callable"))
-            else:
-                # Check if the method signatures are compatible
-                obj_method = getattr(obj, name)
-                proto_sig = inspect.signature(value)
-                obj_sig = inspect.signature(obj_method)
-
-                proto_params = set(proto_sig.parameters)
-                proto_params.discard("self")
-                obj_params = set(obj_sig.parameters)
-                obj_params.discard("self")
-                if not (proto_params <= obj_params):
-                    print(
-                        f"Method {name} incompatible proto: {proto_params} vs. obj: {obj_params}"
-                    )
-                    missing_methods.append((name, "signature_mismatch"))
-                else:
-                    # Check if the method is actually implemented in the class
-                    method_owner = next(
-                        (cls for cls in mro if name in cls.__dict__), None
-                    )
-                    if (
-                        method_owner is None
-                        or method_owner.__name__ == protocol.__name__
-                    ):
-                        missing_methods.append((name, "not_actually_implemented"))
-
-    if missing_methods:
-        raise ValueError(
-            f"Provider `{obj.__provider_id__} ({obj.__provider_spec__.api})` does not implement the following methods:\n{missing_methods}"
-        )

llama_stack/distribution/routers/__init__.py

@@ -1,48 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from typing import Any
-
-from llama_stack.distribution.datatypes import *  # noqa: F403
-from .routing_tables import (
-    MemoryBanksRoutingTable,
-    ModelsRoutingTable,
-    ShieldsRoutingTable,
-)
-
-
-async def get_routing_table_impl(
-    api: Api,
-    impls_by_provider_id: Dict[str, RoutedProtocol],
-    _deps,
-) -> Any:
-    api_to_tables = {
-        "memory_banks": MemoryBanksRoutingTable,
-        "models": ModelsRoutingTable,
-        "shields": ShieldsRoutingTable,
-    }
-    if api.value not in api_to_tables:
-        raise ValueError(f"API {api.value} not found in router map")
-
-    impl = api_to_tables[api.value](impls_by_provider_id)
-    await impl.initialize()
-    return impl
-
-
-async def get_auto_router_impl(api: Api, routing_table: RoutingTable, _deps) -> Any:
-    from .routers import InferenceRouter, MemoryRouter, SafetyRouter
-
-    api_to_routers = {
-        "memory": MemoryRouter,
-        "inference": InferenceRouter,
-        "safety": SafetyRouter,
-    }
-    if api.value not in api_to_routers:
-        raise ValueError(f"API {api.value} not found in router map")
-
-    impl = api_to_routers[api.value](routing_table)
-    await impl.initialize()
-    return impl

llama_stack/distribution/routers/routers.py

@@ -1,158 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from typing import Any, AsyncGenerator, Dict, List
-
-from llama_stack.distribution.datatypes import RoutingTable
-
-from llama_stack.apis.memory import *  # noqa: F403
-from llama_stack.apis.inference import *  # noqa: F403
-from llama_stack.apis.safety import *  # noqa: F403
-
-
-class MemoryRouter(Memory):
-    """Routes to an provider based on the memory bank identifier"""
-
-    def __init__(
-        self,
-        routing_table: RoutingTable,
-    ) -> None:
-        self.routing_table = routing_table
-
-    async def initialize(self) -> None:
-        pass
-
-    async def shutdown(self) -> None:
-        pass
-
-    async def register_memory_bank(self, memory_bank: MemoryBankDef) -> None:
-        await self.routing_table.register_memory_bank(memory_bank)
-
-    async def insert_documents(
-        self,
-        bank_id: str,
-        documents: List[MemoryBankDocument],
-        ttl_seconds: Optional[int] = None,
-    ) -> None:
-        return await self.routing_table.get_provider_impl(bank_id).insert_documents(
-            bank_id, documents, ttl_seconds
-        )
-
-    async def query_documents(
-        self,
-        bank_id: str,
-        query: InterleavedTextMedia,
-        params: Optional[Dict[str, Any]] = None,
-    ) -> QueryDocumentsResponse:
-        return await self.routing_table.get_provider_impl(bank_id).query_documents(
-            bank_id, query, params
-        )
-
-
-class InferenceRouter(Inference):
-    """Routes to an provider based on the model"""
-
-    def __init__(
-        self,
-        routing_table: RoutingTable,
-    ) -> None:
-        self.routing_table = routing_table
-
-    async def initialize(self) -> None:
-        pass
-
-    async def shutdown(self) -> None:
-        pass
-
-    async def register_model(self, model: ModelDef) -> None:
-        await self.routing_table.register_model(model)
-
-    def chat_completion(
-        self,
-        model: str,
-        messages: List[Message],
-        sampling_params: Optional[SamplingParams] = SamplingParams(),
-        tools: Optional[List[ToolDefinition]] = None,
-        tool_choice: Optional[ToolChoice] = ToolChoice.auto,
-        tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json,
-        stream: Optional[bool] = False,
-        logprobs: Optional[LogProbConfig] = None,
-    ) -> AsyncGenerator:
-        params = dict(
-            model=model,
-            messages=messages,
-            sampling_params=sampling_params,
-            tools=tools or [],
-            tool_choice=tool_choice,
-            tool_prompt_format=tool_prompt_format,
-            stream=stream,
-            logprobs=logprobs,
-        )
-        provider = self.routing_table.get_provider_impl(model)
-        if stream:
-            return (chunk async for chunk in provider.chat_completion(**params))
-        else:
-            return provider.chat_completion(**params)
-
-    def completion(
-        self,
-        model: str,
-        content: InterleavedTextMedia,
-        sampling_params: Optional[SamplingParams] = SamplingParams(),
-        stream: Optional[bool] = False,
-        logprobs: Optional[LogProbConfig] = None,
-    ) -> AsyncGenerator:
-        provider = self.routing_table.get_provider_impl(model)
-        params = dict(
-            model=model,
-            content=content,
-            sampling_params=sampling_params,
-            stream=stream,
-            logprobs=logprobs,
-        )
-        if stream:
-            return (chunk async for chunk in provider.completion(**params))
-        else:
-            return provider.completion(**params)
-
-    async def embeddings(
-        self,
-        model: str,
-        contents: List[InterleavedTextMedia],
-    ) -> EmbeddingsResponse:
-        return await self.routing_table.get_provider_impl(model).embeddings(
-            model=model,
-            contents=contents,
-        )
-
-
-class SafetyRouter(Safety):
-    def __init__(
-        self,
-        routing_table: RoutingTable,
-    ) -> None:
-        self.routing_table = routing_table
-
-    async def initialize(self) -> None:
-        pass
-
-    async def shutdown(self) -> None:
-        pass
-
-    async def register_shield(self, shield: ShieldDef) -> None:
-        await self.routing_table.register_shield(shield)
-
-    async def run_shield(
-        self,
-        shield_type: str,
-        messages: List[Message],
-        params: Dict[str, Any] = None,
-    ) -> RunShieldResponse:
-        return await self.routing_table.get_provider_impl(shield_type).run_shield(
-            shield_type=shield_type,
-            messages=messages,
-            params=params,
-        )

llama_stack/distribution/routers/routing_tables.py

@@ -1,173 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from typing import Any, Dict, List, Optional
-
-from llama_models.llama3.api.datatypes import *  # noqa: F403
-
-from llama_stack.apis.models import *  # noqa: F403
-from llama_stack.apis.shields import *  # noqa: F403
-from llama_stack.apis.memory_banks import *  # noqa: F403
-
-from llama_stack.distribution.datatypes import *  # noqa: F403
-
-
-def get_impl_api(p: Any) -> Api:
-    return p.__provider_spec__.api
-
-
-async def register_object_with_provider(obj: RoutableObject, p: Any) -> None:
-    api = get_impl_api(p)
-    if api == Api.inference:
-        await p.register_model(obj)
-    elif api == Api.safety:
-        await p.register_shield(obj)
-    elif api == Api.memory:
-        await p.register_memory_bank(obj)
-
-
-Registry = Dict[str, List[RoutableObjectWithProvider]]
-
-
-# TODO: this routing table maintains state in memory purely. We need to
-# add persistence to it when we add dynamic registration of objects.
-class CommonRoutingTableImpl(RoutingTable):
-    def __init__(
-        self,
-        impls_by_provider_id: Dict[str, RoutedProtocol],
-    ) -> None:
-        self.impls_by_provider_id = impls_by_provider_id
-
-    async def initialize(self) -> None:
-        self.registry: Registry = {}
-
-        def add_objects(objs: List[RoutableObjectWithProvider]) -> None:
-            for obj in objs:
-                if obj.identifier not in self.registry:
-                    self.registry[obj.identifier] = []
-
-                self.registry[obj.identifier].append(obj)
-
-        for pid, p in self.impls_by_provider_id.items():
-            api = get_impl_api(p)
-            if api == Api.inference:
-                p.model_store = self
-                models = await p.list_models()
-                add_objects(
-                    [ModelDefWithProvider(**m.dict(), provider_id=pid) for m in models]
-                )
-
-            elif api == Api.safety:
-                p.shield_store = self
-                shields = await p.list_shields()
-                add_objects(
-                    [
-                        ShieldDefWithProvider(**s.dict(), provider_id=pid)
-                        for s in shields
-                    ]
-                )
-
-            elif api == Api.memory:
-                p.memory_bank_store = self
-                memory_banks = await p.list_memory_banks()
-
-                # do in-memory updates due to pesky Annotated unions
-                for m in memory_banks:
-                    m.provider_id = pid
-
-                add_objects(memory_banks)
-
-    async def shutdown(self) -> None:
-        for p in self.impls_by_provider_id.values():
-            await p.shutdown()
-
-    def get_provider_impl(
-        self, routing_key: str, provider_id: Optional[str] = None
-    ) -> Any:
-        if routing_key not in self.registry:
-            raise ValueError(f"`{routing_key}` not registered")
-
-        objs = self.registry[routing_key]
-        for obj in objs:
-            if not provider_id or provider_id == obj.provider_id:
-                return self.impls_by_provider_id[obj.provider_id]
-
-        raise ValueError(f"Provider not found for `{routing_key}`")
-
-    def get_object_by_identifier(
-        self, identifier: str
-    ) -> Optional[RoutableObjectWithProvider]:
-        objs = self.registry.get(identifier, [])
-        if not objs:
-            return None
-
-        # kind of ill-defined behavior here, but we'll just return the first one
-        return objs[0]
-
-    async def register_object(self, obj: RoutableObjectWithProvider):
-        entries = self.registry.get(obj.identifier, [])
-        for entry in entries:
-            if entry.provider_id == obj.provider_id:
-                print(f"`{obj.identifier}` already registered with `{obj.provider_id}`")
-                return
-
-        if obj.provider_id not in self.impls_by_provider_id:
-            raise ValueError(f"Provider `{obj.provider_id}` not found")
-
-        p = self.impls_by_provider_id[obj.provider_id]
-        await register_object_with_provider(obj, p)
-
-        if obj.identifier not in self.registry:
-            self.registry[obj.identifier] = []
-        self.registry[obj.identifier].append(obj)
-
-        # TODO: persist this to a store
-
-
-class ModelsRoutingTable(CommonRoutingTableImpl, Models):
-    async def list_models(self) -> List[ModelDefWithProvider]:
-        objects = []
-        for objs in self.registry.values():
-            objects.extend(objs)
-        return objects
-
-    async def get_model(self, identifier: str) -> Optional[ModelDefWithProvider]:
-        return self.get_object_by_identifier(identifier)
-
-    async def register_model(self, model: ModelDefWithProvider) -> None:
-        await self.register_object(model)
-
-
-class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):
-    async def list_shields(self) -> List[ShieldDef]:
-        objects = []
-        for objs in self.registry.values():
-            objects.extend(objs)
-        return objects
-
-    async def get_shield(self, shield_type: str) -> Optional[ShieldDefWithProvider]:
-        return self.get_object_by_identifier(shield_type)
-
-    async def register_shield(self, shield: ShieldDefWithProvider) -> None:
-        await self.register_object(shield)
-
-
-class MemoryBanksRoutingTable(CommonRoutingTableImpl, MemoryBanks):
-    async def list_memory_banks(self) -> List[MemoryBankDefWithProvider]:
-        objects = []
-        for objs in self.registry.values():
-            objects.extend(objs)
-        return objects
-
-    async def get_memory_bank(
-        self, identifier: str
-    ) -> Optional[MemoryBankDefWithProvider]:
-        return self.get_object_by_identifier(identifier)
-
-    async def register_memory_bank(
-        self, memory_bank: MemoryBankDefWithProvider
-    ) -> None:
-        await self.register_object(memory_bank)