feat: Support function call requests/responses (#56)

Author: caffeinism

langchain_openai_api_bridge/chat_completion/chat_completion_chunk_choice_adapter.py

@@ -4,33 +4,37 @@
 from langchain_openai_api_bridge.chat_completion.chat_completion_chunk_object_factory import (
     create_chat_completion_chunk_object,
 )
-from langchain_openai_api_bridge.chat_completion.content_adapter import (
-    to_string_content,
-)
-from langchain_openai_api_bridge.core.types.openai import (
-    OpenAIChatCompletionChunkChoice,
-    OpenAIChatCompletionChunkObject,
-    OpenAIChatMessage,
-)
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, Choice, ChoiceDelta, ChoiceDeltaFunctionCall
 
 
 def to_openai_chat_message(
     event: StreamEvent,
     role: str = "assistant",
-) -> OpenAIChatMessage:
-    content = event["data"]["chunk"].content
-    return OpenAIChatMessage(content=to_string_content(content), role=role)
+) -> ChoiceDelta:
+    if getattr(event["data"]["chunk"], "tool_call_chunks", None):
+        function_call = ChoiceDeltaFunctionCall(
+            name=event["data"]["chunk"].tool_call_chunks[0]["name"],
+            arguments=event["data"]["chunk"].tool_call_chunks[0]["args"],
+        )
+    else:
+        function_call = None
+
+    return ChoiceDelta(
+        content=event["data"]["chunk"].content,
+        role=role,
+        function_call=function_call,
+    )
 
 
 def to_openai_chat_completion_chunk_choice(
     event: StreamEvent,
     index: int = 0,
-    role: str = "assistant",
+    role: Optional[str] = None,
     finish_reason: Optional[str] = None,
-) -> OpenAIChatCompletionChunkChoice:
+) -> Choice:
     message = to_openai_chat_message(event, role)
 
-    return OpenAIChatCompletionChunkChoice(
+    return Choice(
         index=index,
         delta=message,
         finish_reason=finish_reason,
@@ -42,9 +46,9 @@ def to_openai_chat_completion_chunk_object(
     id: str = "",
     model: str = "",
     system_fingerprint: Optional[str] = None,
-    role: str = "assistant",
+    role: Optional[str] = None,
     finish_reason: Optional[str] = None,
-) -> OpenAIChatCompletionChunkObject:
+) -> ChatCompletionChunk:
 
     choice1 = to_openai_chat_completion_chunk_choice(
         event, index=0, role=role, finish_reason=finish_reason

langchain_openai_api_bridge/chat_completion/chat_completion_chunk_object_factory.py

@@ -1,19 +1,16 @@
 import time
-from typing import Dict, List, Optional
+from typing import List, Literal, Optional
 
-from langchain_openai_api_bridge.core.types.openai import (
-    OpenAIChatCompletionChunkChoice,
-    OpenAIChatCompletionChunkObject,
-)
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, Choice
 
 
 def create_chat_completion_chunk_object(
     id: str,
     model: str,
     system_fingerprint: Optional[str],
-    choices: List[OpenAIChatCompletionChunkChoice] = [],
-) -> OpenAIChatCompletionChunkObject:
-    return OpenAIChatCompletionChunkObject(
+    choices: List[Choice] = [],
+) -> ChatCompletionChunk:
+    return ChatCompletionChunk(
         id=id,
         object="chat.completion.chunk",
         created=int(time.time()),
@@ -25,18 +22,24 @@ def create_chat_completion_chunk_object(
 
 def create_final_chat_completion_chunk_choice(
     index: int,
-) -> OpenAIChatCompletionChunkChoice:
-    return OpenAIChatCompletionChunkChoice(index=index, delta={}, finish_reason="stop")
+    finish_reason: Literal["stop", "tool_calls"],
+) -> Choice:
+    return Choice(
+        index=index,
+        delta={},
+        finish_reason=finish_reason,
+    )
 
 
 def create_final_chat_completion_chunk_object(
     id: str,
     model: str = "",
     system_fingerprint: Optional[str] = None,
-) -> Dict:
+    finish_reason: Literal["stop", "tool_calls"] = "stop",
+) -> ChatCompletionChunk:
     return create_chat_completion_chunk_object(
         id=id,
         model=model,
         system_fingerprint=system_fingerprint,
-        choices=[create_final_chat_completion_chunk_choice(index=0)],
+        choices=[create_final_chat_completion_chunk_choice(index=0, finish_reason=finish_reason)],
     )

langchain_openai_api_bridge/chat_completion/chat_completion_compatible_api.py

@@ -7,8 +7,8 @@
 from langchain_openai_api_bridge.chat_completion.langchain_stream_adapter import (
     LangchainStreamAdapter,
 )
-from langchain_openai_api_bridge.core.types.openai import OpenAIChatMessage
 from langchain_openai_api_bridge.core.utils.pydantic_async_iterator import ato_dict
+from openai.types.chat import ChatCompletionMessage
 
 
 class ChatCompletionCompatibleAPI:
@@ -39,7 +39,7 @@ def __init__(
         self.agent = agent
         self.event_adapter = event_adapter
 
-    async def astream(self, messages: List[OpenAIChatMessage]) -> AsyncIterator[dict]:
+    async def astream(self, messages: List[ChatCompletionMessage]) -> AsyncIterator[dict]:
         async with self.agent as runnable:
             input = self.__to_input(runnable, messages)
             astream_event = runnable.astream_events(
@@ -51,7 +51,7 @@ async def astream(self, messages: List[OpenAIChatMessage]) -> AsyncIterator[dict
             ):
                 yield it
 
-    async def ainvoke(self, messages: List[OpenAIChatMessage]) -> dict:
+    async def ainvoke(self, messages: List[ChatCompletionMessage]) -> dict:
         async with self.agent as runnable:
             input = self.__to_input(runnable, messages)
             result = await runnable.ainvoke(
@@ -60,16 +60,16 @@ async def ainvoke(self, messages: List[OpenAIChatMessage]) -> dict:
 
         return self.invoke_adapter.to_chat_completion_object(result).model_dump()
 
-    def __to_input(self, runnable: Runnable, messages: List[OpenAIChatMessage]):
+    def __to_input(self, runnable: Runnable, messages: List[ChatCompletionMessage]):
         if isinstance(runnable, CompiledStateGraph):
             return self.__to_react_agent_input(messages)
         else:
             return self.__to_chat_model_input(messages)
 
-    def __to_react_agent_input(self, messages: List[OpenAIChatMessage]):
+    def __to_react_agent_input(self, messages: List[ChatCompletionMessage]):
         return {
-            "messages": [message.model_dump() for message in messages],
+            "messages": [message for message in messages],
         }
 
-    def __to_chat_model_input(self, messages: List[OpenAIChatMessage]):
-        return [message.model_dump() for message in messages]
+    def __to_chat_model_input(self, messages: List[ChatCompletionMessage]):
+        return [message for message in messages]

langchain_openai_api_bridge/chat_completion/chat_completion_object_factory.py

@@ -1,30 +1,26 @@
 import time
 from typing import List, Optional
 
-from langchain_openai_api_bridge.core.types.openai import (
-    OpenAIChatCompletionChoice,
-    OpenAIChatCompletionObject,
-    OpenAIChatCompletionUsage,
-)
+from openai.types.chat.chat_completion import ChatCompletion, Choice, CompletionUsage
 
 
 class ChatCompletionObjectFactory:
     def create(
         id: str,
         model: str,
-        choices: List[OpenAIChatCompletionChoice] = [],
+        choices: List[Choice] = [],
         usage: Optional[
-            OpenAIChatCompletionUsage
-        ] = OpenAIChatCompletionUsage.default(),
+            CompletionUsage
+        ] = CompletionUsage(completion_tokens=-1, prompt_tokens=-1, total_tokens=-1),
         object: str = "chat.completion",
         system_fingerprint: str = "",
         created: int = None,
-    ) -> OpenAIChatCompletionObject:
-        return OpenAIChatCompletionObject(
+    ) -> ChatCompletion:
+        return ChatCompletion(
             id=id,
-            object=object,
-            created=created if created is not None else int(time.time()),
             model=model,
+            created=created or int(time.time()),
+            object=object,
             system_fingerprint=system_fingerprint,
             choices=choices,
             usage=usage,

langchain_openai_api_bridge/chat_completion/langchain_invoke_adapter.py

@@ -1,56 +1,43 @@
+import time
+from langchain_core.messages import BaseMessage
+from langchain_openai.chat_models.base import _convert_message_to_dict
+from openai.types.chat.chat_completion import ChatCompletion, Choice, ChatCompletionMessage
+
 from langchain_openai_api_bridge.chat_completion.chat_completion_object_factory import (
     ChatCompletionObjectFactory,
 )
-from langchain_openai_api_bridge.chat_completion.content_adapter import (
-    to_string_content,
-)
-from langchain_openai_api_bridge.core.role_adapter import to_openai_role
-from langchain_openai_api_bridge.core.types.openai import (
-    OpenAIChatCompletionChoice,
-    OpenAIChatCompletionObject,
-    OpenAIChatMessage,
-)
-from langchain_core.messages import AIMessage
+from langchain_core.runnables.utils import Output
 
 
 class LangchainInvokeAdapter:
     def __init__(self, llm_model: str, system_fingerprint: str = ""):
         self.llm_model = llm_model
         self.system_fingerprint = system_fingerprint
 
-    def to_chat_completion_object(self, invoke_result) -> OpenAIChatCompletionObject:
-        message = self.__create_openai_chat_message(invoke_result)
-        id = self.__get_id(invoke_result)
+    def to_chat_completion_object(self, invoke_result: Output) -> ChatCompletion:
+        invoke_message = invoke_result if isinstance(invoke_result, BaseMessage) else invoke_result["messages"][-1]
+        message = self.__create_openai_chat_message(invoke_message)
+        id = self.__get_id(invoke_message)
 
         return ChatCompletionObjectFactory.create(
             id=id,
             model=self.llm_model,
+            created=int(time.time()),
+            object="chat.completion",
             system_fingerprint=self.system_fingerprint,
             choices=[
-                OpenAIChatCompletionChoice(
+                Choice(
                     index=0,
                     message=message,
-                    finish_reason="stop",
+                    finish_reason="tool_calls" if "tool_calls" in message else "stop",
                 )
-            ],
+            ]
         )
 
-    def __get_id(self, invoke_result):
-        if isinstance(invoke_result, AIMessage):
-            return invoke_result.id
+    def __create_openai_chat_message(self, message: BaseMessage) -> ChatCompletionMessage:
+        message = _convert_message_to_dict(message)
+        message["role"] = "assistant"
+        return message
 
-        last_message = invoke_result["messages"][-1]
-        return last_message.id
-
-    def __create_openai_chat_message(self, invoke_result) -> OpenAIChatMessage:
-        if isinstance(invoke_result, AIMessage):
-            return OpenAIChatMessage(
-                role=to_openai_role(invoke_result.type),
-                content=to_string_content(content=invoke_result.content),
-            )
-
-        last_message = invoke_result["messages"][-1]
-        return OpenAIChatMessage(
-            role=to_openai_role(last_message.type),
-            content=to_string_content(content=last_message.content),
-        )
+    def __get_id(self, message: BaseMessage):
+        return message.id or ""

langchain_openai_api_bridge/chat_completion/langchain_stream_adapter.py

@@ -8,9 +8,7 @@
 from langchain_openai_api_bridge.chat_completion.chat_completion_chunk_object_factory import (
     create_final_chat_completion_chunk_object,
 )
-from langchain_openai_api_bridge.core.types.openai import (
-    OpenAIChatCompletionChunkObject,
-)
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 
 
 class LangchainStreamAdapter:
@@ -23,22 +21,31 @@ async def ato_chat_completion_chunk_stream(
         astream_event: AsyncIterator[StreamEvent],
         id: str = "",
         event_adapter=lambda event: None,
-    ) -> AsyncIterator[OpenAIChatCompletionChunkObject]:
+    ) -> AsyncIterator[ChatCompletionChunk]:
         if id == "":
             id = str(uuid.uuid4())
+
+        is_function_call_prev = is_function_call = False
+        role = "assistant"
         async for event in astream_event:
             custom_event = event_adapter(event)
             event_to_process = custom_event if custom_event is not None else event
             kind = event_to_process["event"]
             if kind == "on_chat_model_stream" or custom_event is not None:
-                yield to_openai_chat_completion_chunk_object(
+                chat_completion_chunk = to_openai_chat_completion_chunk_object(
                     event=event_to_process,
                     id=id,
                     model=self.llm_model,
                     system_fingerprint=self.system_fingerprint,
+                    role=role,
                 )
+                role = None
+                yield chat_completion_chunk
+                is_function_call = is_function_call or any(choice.delta.function_call for choice in chat_completion_chunk.choices)
+            elif kind == "on_chat_model_end":
+                is_function_call_prev, is_function_call = is_function_call, False
 
         stop_chunk = create_final_chat_completion_chunk_object(
-            id=id, model=self.llm_model
+            id=id, model=self.llm_model, finish_reason="tool_calls" if is_function_call_prev else "stop"
         )
         yield stop_chunk

langchain_openai_api_bridge/core/create_agent_dto.py

@@ -1,5 +1,6 @@
 from typing import Optional
 from pydantic import BaseModel
+from openai.types.chat import ChatCompletionToolChoiceOptionParam, ChatCompletionToolParam
 
 
 class CreateAgentDto(BaseModel):
@@ -9,3 +10,5 @@ class CreateAgentDto(BaseModel):
     max_tokens: Optional[int] = None
     assistant_id: Optional[str] = ""
     thread_id: Optional[str] = ""
+    tools: list[ChatCompletionToolParam] = []
+    tool_choice: ChatCompletionToolChoiceOptionParam = "none"

langchain_openai_api_bridge/core/types/openai/__init__.py

@@ -1,19 +1,7 @@
-from .message import OpenAIChatMessage
 from .chat_completion import (
     OpenAIChatCompletionRequest,
-    OpenAIChatCompletionUsage,
-    OpenAIChatCompletionChoice,
-    OpenAIChatCompletionObject,
-    OpenAIChatCompletionChunkChoice,
-    OpenAIChatCompletionChunkObject,
 )
 
 __all__ = [
-    "OpenAIChatMessage",
     "OpenAIChatCompletionRequest",
-    "OpenAIChatCompletionUsage",
-    "OpenAIChatCompletionChoice",
-    "OpenAIChatCompletionObject",
-    "OpenAIChatCompletionChunkChoice",
-    "OpenAIChatCompletionChunkObject",
 ]