Introduce Llama stack distributions (#22)

* Add distribution CLI scaffolding

* More progress towards `llama distribution install`

* getting closer to a distro definition, distro install + configure works

* Distribution server now functioning

* read existing configuration, save enums properly

* Remove inference uvicorn server entrypoint and llama inference CLI command

* updated dependency and client model name

* Improved exception handling

* local imports for faster cli

* undo a typo, add a passthrough distribution

* implement full-passthrough in the server

* add safety adapters, configuration handling, server + clients

* cleanup, moving stuff to common, nuke utils

* Add a Path() wrapper at the earliest place

* fixes

* Bring agentic system api to toolchain

Add adapter dependencies and resolve adapters using a topological sort

* refactor to reduce size of `agentic_system`

* move straggler files and fix some important existing bugs

* ApiSurface -> Api

* refactor a method out

* Adapter -> Provider

* Make each inference provider into its own subdirectory

* installation fixes

* Rename Distribution -> DistributionSpec, simplify RemoteProviders

* dict key instead of attr

* update inference config to take model and not model_dir

* Fix passthrough streaming, send headers properly not part of body :facepalm

* update safety to use model sku ids and not model dirs

* Update cli_reference.md

* minor fixes

* add DistributionConfig, fix a bug in model download

* Make install + start scripts do proper configuration automatically

* Update CLI_reference

* Nuke fp8_requirements, fold fbgemm into common requirements

* Update README, add newline between API surface configurations

* Refactor download functionality out of the Command so can be reused

* Add `llama model download` alias for `llama download`

* Show message about checksum file so users can check themselves

* Simpler intro statements

* get ollama working

* Reduce a bunch of dependencies from toolchain

Some improvements to the distribution install script

* Avoid using `conda run` since it buffers everything

* update dependencies and rely on LLAMA_TOOLCHAIN_DIR for dev purposes

* add validation for configuration input

* resort imports

* make optional subclasses default to yes for configuration

* Remove additional_pip_packages; move deps to providers

* for inline make 8b model the default

* Add scripts to MANIFEST

* allow installing from test.pypi.org

* Fix #2 to help with testing packages

* Must install llama-models at that same version first

* fix PIP_ARGS

---------

Co-authored-by: Hardik Shah <hjshah@fb.com>
Co-authored-by: Hardik Shah <hjshah@meta.com>

llama_toolchain/agentic_system/meta_reference/agentic_system.py

@@ -0,0 +1,142 @@
+# 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 logging
+import os
+import uuid
+from typing import AsyncGenerator, Dict
+
+from llama_toolchain.distribution.datatypes import Api, ProviderSpec
+from llama_toolchain.inference.api import Inference
+from llama_toolchain.inference.api.datatypes import BuiltinTool
+from llama_toolchain.safety.api import Safety
+from llama_toolchain.agentic_system.api.endpoints import *  # noqa
+from llama_toolchain.agentic_system.api import (
+    AgenticSystem,
+    AgenticSystemCreateRequest,
+    AgenticSystemCreateResponse,
+    AgenticSystemSessionCreateRequest,
+    AgenticSystemSessionCreateResponse,
+    AgenticSystemTurnCreateRequest,
+)
+
+from .agent_instance import AgentInstance
+
+from .config import AgenticSystemConfig
+
+from .tools.builtin import (
+    BraveSearchTool,
+    CodeInterpreterTool,
+    PhotogenTool,
+    WolframAlphaTool,
+)
+from .tools.safety import with_safety
+
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+async def get_provider_impl(config: AgenticSystemConfig, deps: Dict[Api, ProviderSpec]):
+    assert isinstance(
+        config, AgenticSystemConfig
+    ), f"Unexpected config type: {type(config)}"
+
+    impl = MetaReferenceAgenticSystemImpl(
+        deps[Api.inference],
+        deps[Api.safety],
+    )
+    await impl.initialize()
+    return impl
+
+
+AGENT_INSTANCES_BY_ID = {}
+
+
+class MetaReferenceAgenticSystemImpl(AgenticSystem):
+    def __init__(self, inference_api: Inference, safety_api: Safety):
+        self.inference_api = inference_api
+        self.safety_api = safety_api
+
+    async def initialize(self) -> None:
+        pass
+
+    async def create_agentic_system(
+        self,
+        request: AgenticSystemCreateRequest,
+    ) -> AgenticSystemCreateResponse:
+        system_id = str(uuid.uuid4())
+
+        builtin_tools = []
+        custom_tool_definitions = []
+        cfg = request.instance_config
+        for dfn in cfg.available_tools:
+            if isinstance(dfn.tool_name, BuiltinTool):
+                if dfn.tool_name == BuiltinTool.wolfram_alpha:
+                    tool = WolframAlphaTool(os.environ.get("WOLFRAM_ALPHA_API_KEY"))
+                elif dfn.tool_name == BuiltinTool.brave_search:
+                    tool = BraveSearchTool(os.environ.get("BRAVE_SEARCH_API_KEY"))
+                elif dfn.tool_name == BuiltinTool.code_interpreter:
+                    tool = CodeInterpreterTool()
+                elif dfn.tool_name == BuiltinTool.photogen:
+                    tool = PhotogenTool(
+                        dump_dir="/tmp/photogen_dump_" + os.environ["USER"],
+                    )
+                else:
+                    raise ValueError(f"Unknown builtin tool: {dfn.tool_name}")
+
+                builtin_tools.append(
+                    with_safety(
+                        tool, self.safety_api, dfn.input_shields, dfn.output_shields
+                    )
+                )
+            else:
+                custom_tool_definitions.append(dfn)
+
+        AGENT_INSTANCES_BY_ID[system_id] = AgentInstance(
+            system_id=system_id,
+            instance_config=request.instance_config,
+            model=request.model,
+            inference_api=self.inference_api,
+            builtin_tools=builtin_tools,
+            custom_tool_definitions=custom_tool_definitions,
+            safety_api=self.safety_api,
+            input_shields=cfg.input_shields,
+            output_shields=cfg.output_shields,
+            prefix_messages=cfg.debug_prefix_messages,
+        )
+
+        return AgenticSystemCreateResponse(
+            system_id=system_id,
+        )
+
+    async def create_agentic_system_session(
+        self,
+        request: AgenticSystemSessionCreateRequest,
+    ) -> AgenticSystemSessionCreateResponse:
+        system_id = request.system_id
+        assert system_id in AGENT_INSTANCES_BY_ID, f"System {system_id} not found"
+        agent = AGENT_INSTANCES_BY_ID[system_id]
+
+        session = agent.create_session(request.session_name)
+        return AgenticSystemSessionCreateResponse(
+            session_id=session.session_id,
+        )
+
+    async def create_agentic_system_turn(
+        self,
+        request: AgenticSystemTurnCreateRequest,
+    ) -> AsyncGenerator:
+        system_id = request.system_id
+        assert system_id in AGENT_INSTANCES_BY_ID, f"System {system_id} not found"
+        agent = AGENT_INSTANCES_BY_ID[system_id]
+
+        assert (
+            request.session_id in agent.sessions
+        ), f"Session {request.session_id} not found"
+        async for event in agent.create_and_execute_turn(request):
+            yield event