rebase on top of registry

llama_stack/providers/utils/inference/model_registry.py

@@ -0,0 +1,35 @@
+# 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 Dict
+
+from llama_models.sku_list import resolve_model
+
+from llama_stack.apis.models import *  # noqa: F403
+
+
+class ModelRegistryHelper:
+
+    def __init__(self, stack_to_provider_models_map: Dict[str, str]):
+        self.stack_to_provider_models_map = stack_to_provider_models_map
+
+    def map_to_provider_model(self, identifier: str) -> str:
+        model = resolve_model(identifier)
+        if not model:
+            raise ValueError(f"Unknown model: `{identifier}`")
+
+        if identifier not in self.stack_to_provider_models_map:
+            raise ValueError(
+                f"Model {identifier} not found in map {self.stack_to_provider_models_map}"
+            )
+
+        return self.stack_to_provider_models_map[identifier]
+
+    async def register_model(self, model: ModelDef) -> None:
+        if model.identifier not in self.stack_to_provider_models_map:
+            raise ValueError(
+                f"Unsupported model {model.identifier}. Supported models: {self.stack_to_provider_models_map.keys()}"
+            )

llama_stack/providers/utils/inference/openai_compat.py

@@ -0,0 +1,189 @@
+# 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 AsyncGenerator, Optional
+
+from llama_models.llama3.api.chat_format import ChatFormat
+
+from llama_models.llama3.api.datatypes import StopReason
+
+from llama_stack.apis.inference import *  # noqa: F403
+
+from pydantic import BaseModel
+
+
+class OpenAICompatCompletionChoiceDelta(BaseModel):
+    content: str
+
+
+class OpenAICompatCompletionChoice(BaseModel):
+    finish_reason: Optional[str] = None
+    text: Optional[str] = None
+    delta: Optional[OpenAICompatCompletionChoiceDelta] = None
+
+
+class OpenAICompatCompletionResponse(BaseModel):
+    choices: List[OpenAICompatCompletionChoice]
+
+
+def get_sampling_options(request: ChatCompletionRequest) -> dict:
+    options = {}
+    if params := request.sampling_params:
+        for attr in {"temperature", "top_p", "top_k", "max_tokens"}:
+            if getattr(params, attr):
+                options[attr] = getattr(params, attr)
+
+        if params.repetition_penalty is not None and params.repetition_penalty != 1.0:
+            options["repeat_penalty"] = params.repetition_penalty
+
+    return options
+
+
+def text_from_choice(choice) -> str:
+    if hasattr(choice, "delta") and choice.delta:
+        return choice.delta.content
+
+    return choice.text
+
+
+def process_chat_completion_response(
+    request: ChatCompletionRequest,
+    response: OpenAICompatCompletionResponse,
+    formatter: ChatFormat,
+) -> ChatCompletionResponse:
+    choice = response.choices[0]
+
+    stop_reason = None
+    if reason := choice.finish_reason:
+        if reason in ["stop", "eos"]:
+            stop_reason = StopReason.end_of_turn
+        elif reason == "eom":
+            stop_reason = StopReason.end_of_message
+        elif reason == "length":
+            stop_reason = StopReason.out_of_tokens
+
+    if stop_reason is None:
+        stop_reason = StopReason.out_of_tokens
+
+    completion_message = formatter.decode_assistant_message_from_content(
+        text_from_choice(choice), stop_reason
+    )
+    return ChatCompletionResponse(
+        completion_message=completion_message,
+        logprobs=None,
+    )
+
+
+async def process_chat_completion_stream_response(
+    request: ChatCompletionRequest,
+    stream: AsyncGenerator[OpenAICompatCompletionResponse, None],
+    formatter: ChatFormat,
+) -> AsyncGenerator:
+    yield ChatCompletionResponseStreamChunk(
+        event=ChatCompletionResponseEvent(
+            event_type=ChatCompletionResponseEventType.start,
+            delta="",
+        )
+    )
+
+    buffer = ""
+    ipython = False
+    stop_reason = None
+
+    async for chunk in stream:
+        choice = chunk.choices[0]
+        finish_reason = choice.finish_reason
+
+        if finish_reason:
+            if stop_reason is None and finish_reason in ["stop", "eos", "eos_token"]:
+                stop_reason = StopReason.end_of_turn
+            elif stop_reason is None and finish_reason == "length":
+                stop_reason = StopReason.out_of_tokens
+            break
+
+        text = text_from_choice(choice)
+        # check if its a tool call ( aka starts with <|python_tag|> )
+        if not ipython and text.startswith("<|python_tag|>"):
+            ipython = True
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=ToolCallDelta(
+                        content="",
+                        parse_status=ToolCallParseStatus.started,
+                    ),
+                )
+            )
+            buffer += text
+            continue
+
+        if text == "<|eot_id|>":
+            stop_reason = StopReason.end_of_turn
+            text = ""
+            continue
+        elif text == "<|eom_id|>":
+            stop_reason = StopReason.end_of_message
+            text = ""
+            continue
+
+        if ipython:
+            buffer += text
+            delta = ToolCallDelta(
+                content=text,
+                parse_status=ToolCallParseStatus.in_progress,
+            )
+
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=delta,
+                    stop_reason=stop_reason,
+                )
+            )
+        else:
+            buffer += text
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=text,
+                    stop_reason=stop_reason,
+                )
+            )
+
+    # parse tool calls and report errors
+    message = formatter.decode_assistant_message_from_content(buffer, stop_reason)
+    parsed_tool_calls = len(message.tool_calls) > 0
+    if ipython and not parsed_tool_calls:
+        yield ChatCompletionResponseStreamChunk(
+            event=ChatCompletionResponseEvent(
+                event_type=ChatCompletionResponseEventType.progress,
+                delta=ToolCallDelta(
+                    content="",
+                    parse_status=ToolCallParseStatus.failure,
+                ),
+                stop_reason=stop_reason,
+            )
+        )
+
+    for tool_call in message.tool_calls:
+        yield ChatCompletionResponseStreamChunk(
+            event=ChatCompletionResponseEvent(
+                event_type=ChatCompletionResponseEventType.progress,
+                delta=ToolCallDelta(
+                    content=tool_call,
+                    parse_status=ToolCallParseStatus.success,
+                ),
+                stop_reason=stop_reason,
+            )
+        )
+
+    yield ChatCompletionResponseStreamChunk(
+        event=ChatCompletionResponseEvent(
+            event_type=ChatCompletionResponseEventType.complete,
+            delta="",
+            stop_reason=stop_reason,
+        )
+    )

llama_stack/providers/utils/inference/prompt_adapter.py

@@ -3,7 +3,11 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+from typing import Tuple
+
+from llama_models.llama3.api.chat_format import ChatFormat
 from termcolor import cprint
+
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.inference import *  # noqa: F403
 from llama_models.datatypes import ModelFamily
@@ -19,7 +23,28 @@ from llama_models.sku_list import resolve_model
 from llama_stack.providers.utils.inference import supported_inference_models
 
 
-def augment_messages_for_tools(request: ChatCompletionRequest) -> List[Message]:
+def chat_completion_request_to_prompt(
+    request: ChatCompletionRequest, formatter: ChatFormat
+) -> str:
+    messages = chat_completion_request_to_messages(request)
+    model_input = formatter.encode_dialog_prompt(messages)
+    return formatter.tokenizer.decode(model_input.tokens)
+
+
+def chat_completion_request_to_model_input_info(
+    request: ChatCompletionRequest, formatter: ChatFormat
+) -> Tuple[str, int]:
+    messages = chat_completion_request_to_messages(request)
+    model_input = formatter.encode_dialog_prompt(messages)
+    return (
+        formatter.tokenizer.decode(model_input.tokens),
+        len(model_input.tokens),
+    )
+
+
+def chat_completion_request_to_messages(
+    request: ChatCompletionRequest,
+) -> List[Message]:
     """Reads chat completion request and augments the messages to handle tools.
     For eg. for llama_3_1, add system message with the appropriate tools or
     add user messsage for custom tools, etc.
@@ -48,7 +73,6 @@ def augment_messages_for_tools(request: ChatCompletionRequest) -> List[Message]:
 def augment_messages_for_tools_llama_3_1(
     request: ChatCompletionRequest,
 ) -> List[Message]:
-
     assert request.tool_choice == ToolChoice.auto, "Only `ToolChoice.auto` supported"
 
     existing_messages = request.messages

llama_stack/providers/utils/inference/routable.py

@@ -1,36 +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 Dict, List
-
-from llama_models.sku_list import resolve_model
-
-from llama_stack.distribution.datatypes import RoutableProvider
-
-
-class RoutableProviderForModels(RoutableProvider):
-
-    def __init__(self, stack_to_provider_models_map: Dict[str, str]):
-        self.stack_to_provider_models_map = stack_to_provider_models_map
-
-    async def validate_routing_keys(self, routing_keys: List[str]):
-        for routing_key in routing_keys:
-            if routing_key not in self.stack_to_provider_models_map:
-                raise ValueError(
-                    f"Routing key {routing_key} not found in map {self.stack_to_provider_models_map}"
-                )
-
-    def map_to_provider_model(self, routing_key: str) -> str:
-        model = resolve_model(routing_key)
-        if not model:
-            raise ValueError(f"Unknown model: `{routing_key}`")
-
-        if routing_key not in self.stack_to_provider_models_map:
-            raise ValueError(
-                f"Model {routing_key} not found in map {self.stack_to_provider_models_map}"
-            )
-
-        return self.stack_to_provider_models_map[routing_key]

llama_stack/providers/utils/kvstore/config.py

@@ -31,6 +31,10 @@ class RedisKVStoreConfig(CommonConfig):
     host: str = "localhost"
     port: int = 6379
 
+    @property
+    def url(self) -> str:
+        return f"redis://{self.host}:{self.port}"
+
 
 class SqliteKVStoreConfig(CommonConfig):
     type: Literal[KVStoreType.sqlite.value] = KVStoreType.sqlite.value

llama_stack/providers/utils/memory/vector_store.py

@@ -146,22 +146,22 @@ class EmbeddingIndex(ABC):
 
 @dataclass
 class BankWithIndex:
-    bank: MemoryBank
+    bank: MemoryBankDef
     index: EmbeddingIndex
 
     async def insert_documents(
         self,
         documents: List[MemoryBankDocument],
     ) -> None:
-        model = get_embedding_model(self.bank.config.embedding_model)
+        model = get_embedding_model(self.bank.embedding_model)
         for doc in documents:
             content = await content_from_doc(doc)
             chunks = make_overlapped_chunks(
                 doc.document_id,
                 content,
-                self.bank.config.chunk_size_in_tokens,
-                self.bank.config.overlap_size_in_tokens
-                or (self.bank.config.chunk_size_in_tokens // 4),
+                self.bank.chunk_size_in_tokens,
+                self.bank.overlap_size_in_tokens
+                or (self.bank.chunk_size_in_tokens // 4),
             )
             if not chunks:
                 continue
@@ -189,6 +189,6 @@ class BankWithIndex:
         else:
             query_str = _process(query)
 
-        model = get_embedding_model(self.bank.config.embedding_model)
+        model = get_embedding_model(self.bank.embedding_model)
         query_vector = model.encode([query_str])[0].astype(np.float32)
         return await self.index.query(query_vector, k)