phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-30 07:39:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

routers wip

Browse source
This commit is contained in:
Xi Yan 2024-09-19 08:32:47 -07:00
parent ef0e717bd0
commit f3ff3a3001
4 changed files with 246 additions and 155 deletions
Download patch file Download diff file Expand all files Collapse all files

311
llama_stack/providers/adapters/inference/ollama/ollama.py
View file

@ -17,6 +17,7 @@ from ollama import AsyncClient
from llama_stack.apis.inference import * # noqa: F403 from llama_stack.apis.inference import * # noqa: F403
from llama_stack.providers.utils.inference.prepare_messages import prepare_messages from llama_stack.providers.utils.inference.prepare_messages import prepare_messages
from termcolor import cprint
# TODO: Eventually this will move to the llama cli model list command # TODO: Eventually this will move to the llama cli model list command
# mapping of Model SKUs to ollama models # mapping of Model SKUs to ollama models
@ -38,12 +39,13 @@ class OllamaInferenceAdapter(Inference):
return AsyncClient(host=self.url) return AsyncClient(host=self.url)
async def initialize(self) -> None: async def initialize(self) -> None:
try: pass
await self.client.ps() # try:
except httpx.ConnectError as e: # await self.client.ps()
raise RuntimeError( # except httpx.ConnectError as e:
"Ollama Server is not running, start it using `ollama serve` in a separate terminal" # raise RuntimeError(
) from e # "Ollama Server is not running, start it using `ollama serve` in a separate terminal"
# ) from e
async def shutdown(self) -> None: async def shutdown(self) -> None:
pass pass
@ -96,166 +98,167 @@ class OllamaInferenceAdapter(Inference):
stream: Optional[bool] = False, stream: Optional[bool] = False,
logprobs: Optional[LogProbConfig] = None, logprobs: Optional[LogProbConfig] = None,
) -> AsyncGenerator: ) -> AsyncGenerator:
request = ChatCompletionRequest( cprint("!! calling remote ollama !!", "red")
model=model, # request = ChatCompletionRequest(
messages=messages, # model=model,
sampling_params=sampling_params, # messages=messages,
tools=tools or [], # sampling_params=sampling_params,
tool_choice=tool_choice, # tools=tools or [],
tool_prompt_format=tool_prompt_format, # tool_choice=tool_choice,
stream=stream, # tool_prompt_format=tool_prompt_format,
logprobs=logprobs, # stream=stream,
) # logprobs=logprobs,
# )
messages = prepare_messages(request) # messages = prepare_messages(request)
# accumulate sampling params and other options to pass to ollama # # accumulate sampling params and other options to pass to ollama
options = self.get_ollama_chat_options(request) # options = self.get_ollama_chat_options(request)
ollama_model = self.resolve_ollama_model(request.model) # ollama_model = self.resolve_ollama_model(request.model)
res = await self.client.ps() # res = await self.client.ps()
need_model_pull = True # need_model_pull = True
for r in res["models"]: # for r in res["models"]:
if ollama_model == r["model"]: # if ollama_model == r["model"]:
need_model_pull = False # need_model_pull = False
break # break
if need_model_pull: # if need_model_pull:
print(f"Pulling model: {ollama_model}") # print(f"Pulling model: {ollama_model}")
status = await self.client.pull(ollama_model) # status = await self.client.pull(ollama_model)
assert ( # assert (
status["status"] == "success" # status["status"] == "success"
), f"Failed to pull model {self.model} in ollama" # ), f"Failed to pull model {self.model} in ollama"
if not request.stream: # if not request.stream:
r = await self.client.chat( # r = await self.client.chat(
model=ollama_model, # model=ollama_model,
messages=self._messages_to_ollama_messages(messages), # messages=self._messages_to_ollama_messages(messages),
stream=False, # stream=False,
options=options, # options=options,
) # )
stop_reason = None # stop_reason = None
if r["done"]: # if r["done"]:
if r["done_reason"] == "stop": # if r["done_reason"] == "stop":
stop_reason = StopReason.end_of_turn # stop_reason = StopReason.end_of_turn
elif r["done_reason"] == "length": # elif r["done_reason"] == "length":
stop_reason = StopReason.out_of_tokens # stop_reason = StopReason.out_of_tokens
completion_message = self.formatter.decode_assistant_message_from_content( # completion_message = self.formatter.decode_assistant_message_from_content(
r["message"]["content"], stop_reason # r["message"]["content"], stop_reason
) # )
yield ChatCompletionResponse( # yield ChatCompletionResponse(
completion_message=completion_message, # completion_message=completion_message,
logprobs=None, # logprobs=None,
) # )
else: # else:
yield ChatCompletionResponseStreamChunk( # yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( # event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.start, # event_type=ChatCompletionResponseEventType.start,
delta="", # delta="",
) # )
) # )
stream = await self.client.chat( # stream = await self.client.chat(
model=ollama_model, # model=ollama_model,
messages=self._messages_to_ollama_messages(messages), # messages=self._messages_to_ollama_messages(messages),
stream=True, # stream=True,
options=options, # options=options,
) # )
buffer = "" # buffer = ""
ipython = False # ipython = False
stop_reason = None # stop_reason = None
async for chunk in stream: # async for chunk in stream:
if chunk["done"]: # if chunk["done"]:
if stop_reason is None and chunk["done_reason"] == "stop": # if stop_reason is None and chunk["done_reason"] == "stop":
stop_reason = StopReason.end_of_turn # stop_reason = StopReason.end_of_turn
elif stop_reason is None and chunk["done_reason"] == "length": # elif stop_reason is None and chunk["done_reason"] == "length":
stop_reason = StopReason.out_of_tokens # stop_reason = StopReason.out_of_tokens
break # break
text = chunk["message"]["content"] # text = chunk["message"]["content"]
# check if its a tool call ( aka starts with <|python_tag|> ) # # check if its a tool call ( aka starts with <|python_tag|> )
if not ipython and text.startswith("<|python_tag|>"): # if not ipython and text.startswith("<|python_tag|>"):
ipython = True # ipython = True
yield ChatCompletionResponseStreamChunk( # yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( # event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.progress, # event_type=ChatCompletionResponseEventType.progress,
delta=ToolCallDelta( # delta=ToolCallDelta(
content="", # content="",
parse_status=ToolCallParseStatus.started, # parse_status=ToolCallParseStatus.started,
), # ),
) # )
) # )
buffer += text # buffer += text
continue # continue
if ipython: # if ipython:
if text == "<|eot_id|>": # if text == "<|eot_id|>":
stop_reason = StopReason.end_of_turn # stop_reason = StopReason.end_of_turn
text = "" # text = ""
continue # continue
elif text == "<|eom_id|>": # elif text == "<|eom_id|>":
stop_reason = StopReason.end_of_message # stop_reason = StopReason.end_of_message
text = "" # text = ""
continue # continue
buffer += text # buffer += text
delta = ToolCallDelta( # delta = ToolCallDelta(
content=text, # content=text,
parse_status=ToolCallParseStatus.in_progress, # parse_status=ToolCallParseStatus.in_progress,
) # )
yield ChatCompletionResponseStreamChunk( # yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( # event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.progress, # event_type=ChatCompletionResponseEventType.progress,
delta=delta, # delta=delta,
stop_reason=stop_reason, # stop_reason=stop_reason,
) # )
) # )
else: # else:
buffer += text # buffer += text
yield ChatCompletionResponseStreamChunk( # yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( # event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.progress, # event_type=ChatCompletionResponseEventType.progress,
delta=text, # delta=text,
stop_reason=stop_reason, # stop_reason=stop_reason,
) # )
) # )
# parse tool calls and report errors # # parse tool calls and report errors
message = self.formatter.decode_assistant_message_from_content( # message = self.formatter.decode_assistant_message_from_content(
buffer, stop_reason # buffer, stop_reason
) # )
parsed_tool_calls = len(message.tool_calls) > 0 # parsed_tool_calls = len(message.tool_calls) > 0
if ipython and not parsed_tool_calls: # if ipython and not parsed_tool_calls:
yield ChatCompletionResponseStreamChunk( # yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( # event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.progress, # event_type=ChatCompletionResponseEventType.progress,
delta=ToolCallDelta( # delta=ToolCallDelta(
content="", # content="",
parse_status=ToolCallParseStatus.failure, # parse_status=ToolCallParseStatus.failure,
), # ),
stop_reason=stop_reason, # stop_reason=stop_reason,
) # )
) # )
for tool_call in message.tool_calls: # for tool_call in message.tool_calls:
yield ChatCompletionResponseStreamChunk( # yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( # event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.progress, # event_type=ChatCompletionResponseEventType.progress,
delta=ToolCallDelta( # delta=ToolCallDelta(
content=tool_call, # content=tool_call,
parse_status=ToolCallParseStatus.success, # parse_status=ToolCallParseStatus.success,
), # ),
stop_reason=stop_reason, # stop_reason=stop_reason,
) # )
) # )
yield ChatCompletionResponseStreamChunk( # yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( # event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.complete, # event_type=ChatCompletionResponseEventType.complete,
delta="", # delta="",
stop_reason=stop_reason, # stop_reason=stop_reason,
) # )
) # )

25
llama_stack/providers/impls/meta_reference/models/models.py
View file

@ -16,6 +16,9 @@ from llama_models.datatypes import CoreModelId, Model
from llama_models.sku_list import resolve_model from llama_models.sku_list import resolve_model
from llama_stack.apis.inference import Inference from llama_stack.apis.inference import Inference
from llama_stack.apis.safety import Safety from llama_stack.apis.safety import Safety
from llama_stack.providers.adapters.inference.ollama.ollama import (
OllamaInferenceAdapter,
)
from llama_stack.providers.impls.meta_reference.inference.inference import ( from llama_stack.providers.impls.meta_reference.inference.inference import (
MetaReferenceInferenceImpl, MetaReferenceInferenceImpl,
@ -23,6 +26,7 @@ from llama_stack.providers.impls.meta_reference.inference.inference import (
from llama_stack.providers.impls.meta_reference.safety.safety import ( from llama_stack.providers.impls.meta_reference.safety.safety import (
MetaReferenceSafetyImpl, MetaReferenceSafetyImpl,
) )
from llama_stack.providers.routers.inference.inference import InferenceRouterImpl
from .config import MetaReferenceImplConfig from .config import MetaReferenceImplConfig
@ -39,7 +43,7 @@ class MetaReferenceModelsImpl(Models):
self.safety_api = safety_api self.safety_api = safety_api
self.models_list = [] self.models_list = []
model = get_model_id_from_api(self.inference_api) # model = get_model_id_from_api(self.inference_api)
# TODO, make the inference route provider and use router provider to do the lookup dynamically # TODO, make the inference route provider and use router provider to do the lookup dynamically
if isinstance( if isinstance(
@ -56,6 +60,25 @@ class MetaReferenceModelsImpl(Models):
) )
) )
if isinstance(
self.inference_api,
OllamaInferenceAdapter,
):
self.models_list.append(
ModelSpec(
providers_spec={
"inference": [{"provider_type": "remote::ollama"}],
},
)
)
if isinstance(
self.inference_api,
InferenceRouterImpl,
):
print("Found router")
print(self.inference_api.providers)
if isinstance( if isinstance(
self.safety_api, self.safety_api,
MetaReferenceSafetyImpl, MetaReferenceSafetyImpl,

17
llama_stack/providers/routers/inference/__init__.py Normal file
View file

@ -0,0 +1,17 @@
# 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, List, Tuple
from llama_stack.distribution.datatypes import Api
async def get_router_impl(inner_impls: List[Tuple[str, Any]], deps: List[Api]):
from .inference import InferenceRouterImpl
impl = InferenceRouterImpl(inner_impls, deps)
await impl.initialize()
return impl

48
llama_stack/providers/routers/inference/inference.py Normal file
View file

@ -0,0 +1,48 @@
# 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, Tuple
from llama_stack.distribution.datatypes import Api
from llama_stack.apis.inference import * # noqa: F403
from llama_stack.providers.registry.inference import available_providers
class InferenceRouterImpl(Inference):
"""Routes to an provider based on the memory bank type"""
def __init__(
self,
inner_impls: List[Tuple[str, Any]],
deps: List[Api],
) -> None:
self.inner_impls = inner_impls
self.deps = deps
self.providers = {}
for routing_key, provider_impl in inner_impls:
self.providers[routing_key] = provider_impl
async def initialize(self) -> None:
pass
async def shutdown(self) -> None:
for p in self.providers.values():
await p.shutdown()
async def chat_completion(
self,
model: str,
messages: List[Message],
sampling_params: Optional[SamplingParams] = SamplingParams(),
# zero-shot tool definitions as input to the model
tools: Optional[List[ToolDefinition]] = list,
tool_choice: Optional[ToolChoice] = ToolChoice.auto,
tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json,
stream: Optional[bool] = False,
logprobs: Optional[LogProbConfig] = None,
) -> AsyncGenerator:
print("router chat_completion")