import json from typing import Optional from urllib.parse import urlparse from fastapi import HTTPException from lmos_openai_types import ( ChatCompletionMessageToolCall, ChatCompletionRequestMessage, CreateChatCompletionRequest, CreateChatCompletionStreamResponse, Function1, ) from .utils import call_tool, chat_completion_add_tools from mcp_bridge.models import SSEData from .genericHttpxClient import get_client from mcp_bridge.mcp_clients.McpClientManager import ClientManager from mcp_bridge.tool_mappers import mcp2openai from mcp_bridge.config import config from loguru import logger from httpx_sse import aconnect_sse from sse_starlette.sse import EventSourceResponse, ServerSentEvent async def streaming_chat_completions(request: CreateChatCompletionRequest, api_key: str = None,): # raise NotImplementedError("Streaming Chat Completion is not supported") try: return EventSourceResponse( content=chat_completions(request, api_key), media_type="text/event-stream", headers={"Cache-Control": "no-cache"}, ) except Exception as e: logger.error(e) async def chat_completions(request: CreateChatCompletionRequest, api_key: str = None,): """performs a chat completion using the inference server""" client = get_client(api_key) request.stream = True request = await chat_completion_add_tools(request) fully_done = False while not fully_done: # json_data = request.model_dump_json( # exclude_defaults=True, exclude_none=True, exclude_unset=True # ) json_data = json.dumps(request.model_dump( exclude_defaults=True, exclude_none=True, exclude_unset=True )) # logger.debug(json_data) last: Optional[CreateChatCompletionStreamResponse] = None # last message tool_call_name: str = "" tool_call_json: str = "" should_forward: bool = True response_content: str = "" tool_call_id: str = "" async with aconnect_sse( client, "post", "/chat/completions", content=json_data ) as event_source: # check if the content type is correct because the aiter_sse method # will raise an exception if the content type is not correct if "Content-Type" in event_source.response.headers: content_type = event_source.response.headers["Content-Type"] if "text/event-stream" not in content_type: logger.error(f"Unexpected Content-Type: {content_type}") error_data = await event_source.response.aread() logger.error(f"Request URL: {event_source.response.url}") logger.error(f"Request Data: {json_data}") logger.error(f"Response Status: {event_source.response.status_code}") logger.error(f"Response Data: {error_data.decode(event_source.response.encoding or 'utf-8')}") raise HTTPException(status_code=500, detail="Unexpected Content-Type") # iterate over the SSE stream async for sse in event_source.aiter_sse(): event = sse.event data = sse.data id = sse.id retry = sse.retry logger.debug( f"event: {event},\ndata: {data},\nid: {id},\nretry: {retry}" ) # handle if the SSE stream is done if data == "[DONE]": logger.debug("inference serverstream done") break # for some reason openrouter uses uppercase for finish_reason try: data['choices'][0]['finish_reason'] = data['choices'][0]['finish_reason'].lower() # type: ignore except Exception as e: logger.debug(f"failed to lowercase finish_reason: {e}") try: parsed_data = CreateChatCompletionStreamResponse.model_validate_json( data ) except Exception as e: logger.debug(data) raise e # add the delta to the response content content = parsed_data.choices[0].delta.content content = content if content is not None else "" response_content += content # handle stop reasons if parsed_data.choices[0].finish_reason is not None: if parsed_data.choices[0].finish_reason.value in [ "stop", "length", ]: fully_done = True else: should_forward = False # this manages the incoming tool call schema # most of this is assertions to please mypy if parsed_data.choices[0].delta.tool_calls is not None: should_forward = False assert ( parsed_data.choices[0].delta.tool_calls[0].function is not None ) name = parsed_data.choices[0].delta.tool_calls[0].function.name name = name if name is not None else "" tool_call_name = name if tool_call_name == "" else tool_call_name call_id = parsed_data.choices[0].delta.tool_calls[0].id call_id = call_id if call_id is not None else "" tool_call_id = id if tool_call_id == "" else tool_call_id arg = parsed_data.choices[0].delta.tool_calls[0].function.arguments tool_call_json += arg if arg is not None else "" # forward SSE messages to the client logger.debug(f"{should_forward=}") if should_forward: # we do not want to forward tool call json to the client logger.debug("forwarding message") yield SSEData.model_validate_json(sse.data).model_dump_json() # save the last message last = parsed_data # ideally we should check this properly assert last is not None assert last.choices[0].finish_reason is not None if last.choices[0].finish_reason.value in ["stop", "length"]: logger.debug("no tool calls found") fully_done = True continue logger.debug("tool calls found") logger.info( f"{tool_call_name=} {tool_call_json=}" ) # this should not be error but its easier to debug # add received message to the history msg = ChatCompletionRequestMessage( role="assistant", content=response_content, tool_calls=[ ChatCompletionMessageToolCall( id=tool_call_id, type="function", function=Function1(name=tool_call_name, arguments=tool_call_json), ) ], ) # type: ignore request.messages.append(msg) #### MOST OF THIS IS COPY PASTED FROM CHAT_COMPLETIONS # FIXME: this can probably be done in parallel using asyncio gather tool_call_result = await call_tool(tool_call_name, tool_call_json) if tool_call_result is None: continue logger.debug( f"tool call result for {tool_call_name}: {tool_call_result.model_dump()}" ) logger.debug(f"tool call result content: {tool_call_result.content}") # Check if we're using Anthropic API def is_anthropic_endpoint(base_url: str) -> bool: """Detect if we're using Anthropic's API""" parsed_url = urlparse(base_url) return "anthropic.com" in parsed_url.netloc if is_anthropic_endpoint(config.inference_server.base_url): # For Anthropic, use a simple string content text_parts = [part.text for part in filter(lambda x: x.type == "text", tool_call_result.content)] content_str = " ".join(text_parts) if text_parts else "the tool call result is empty" request.messages.append( ChatCompletionRequestMessage.model_validate( { "role": "tool", "content": content_str, # Simple string for Anthropic "tool_call_id": tool_call_id, } ) ) else: # For other providers, keep the rich format tools_content = [ {"type": "text", "text": part.text} for part in filter(lambda x: x.type == "text", tool_call_result.content) ] if len(tools_content) == 0: tools_content = [{"type": "text", "text": "the tool call result is empty"}] request.messages.append( ChatCompletionRequestMessage.model_validate( { "role": "tool", "content": str(tools_content), "tool_call_id": tool_call_id, } ) ) logger.debug("sending next iteration of chat completion request") # when done, send the final event logger.debug("sending final event") yield ServerSentEvent(event="message", data="[DONE]", id=None, retry=None)