feat: support non react-agent for chat completion api (#33)

Author: samuelint

langchain_openai_api_bridge/chat_completion/chat_completion_compatible_api.py

@@ -1,5 +1,6 @@
 from typing import AsyncIterator, List, Optional
 from langchain_core.runnables import Runnable
+from langgraph.graph.graph import CompiledGraph
 from langchain_openai_api_bridge.chat_completion.langchain_invoke_adapter import (
     LangchainInvokeAdapter,
 )
@@ -43,13 +44,24 @@ def astream(self, messages: List[OpenAIChatMessage]) -> AsyncIterator[dict]:
         )
 
     def invoke(self, messages: List[OpenAIChatMessage]) -> dict:
+        input = self.__to_input(messages)
+
         result = self.agent.invoke(
-            input=self.__to_input(messages),
+            input=input,
         )
 
         return self.invoke_adapter.to_chat_completion_object(result).dict()
 
     def __to_input(self, messages: List[OpenAIChatMessage]):
+        if isinstance(self.agent, CompiledGraph):
+            return self.__to_react_agent_input(messages)
+        else:
+            return self.__to_chat_model_input(messages)
+
+    def __to_react_agent_input(self, messages: List[OpenAIChatMessage]):
         return {
             "messages": [message.dict() for message in messages],
         }
+
+    def __to_chat_model_input(self, messages: List[OpenAIChatMessage]):
+        return [message.dict() for message in messages]

langchain_openai_api_bridge/chat_completion/langchain_invoke_adapter.py

@@ -10,6 +10,7 @@
     OpenAIChatCompletionObject,
     OpenAIChatMessage,
 )
+from langchain_core.messages import AIMessage
 
 
 class LangchainInvokeAdapter:
@@ -18,20 +19,38 @@ def __init__(self, llm_model: str, system_fingerprint: str = ""):
         self.system_fingerprint = system_fingerprint
 
     def to_chat_completion_object(self, invoke_result) -> OpenAIChatCompletionObject:
-        last_message = invoke_result["messages"][-1]
+        message = self.__create_openai_chat_message(invoke_result)
+        id = self.__get_id(invoke_result)
 
         return ChatCompletionObjectFactory.create(
-            id=last_message.id,
+            id=id,
             model=self.llm_model,
             system_fingerprint=self.system_fingerprint,
             choices=[
                 OpenAIChatCompletionChoice(
                     index=0,
-                    message=OpenAIChatMessage(
-                        role=to_openai_role(last_message.type),
-                        content=to_string_content(content=last_message.content),
-                    ),
+                    message=message,
                     finish_reason="stop",
                 )
             ],
         )
+
+    def __get_id(self, invoke_result):
+        if isinstance(invoke_result, AIMessage):
+            return invoke_result.id
+
+        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),
+        )

tests/test_functional/fastapi_chat_completion_openai/my_openai_agent_factory.py

@@ -1,18 +1,10 @@
 from langchain_openai_api_bridge.core.base_agent_factory import BaseAgentFactory
 from langchain_core.runnables import Runnable
-from langchain_core.tools import tool
-from langgraph.prebuilt import create_react_agent
 from langchain_openai import ChatOpenAI
 
 from langchain_openai_api_bridge.core.create_agent_dto import CreateAgentDto
 
 
-@tool
-def magic_number_tool(input: int) -> int:
-    """Applies a magic function to an input."""
-    return input + 2
-
-
 class MyOpenAIAgentFactory(BaseAgentFactory):
 
     def create_agent(self, dto: CreateAgentDto) -> Runnable:
@@ -23,8 +15,4 @@ def create_agent(self, dto: CreateAgentDto) -> Runnable:
             temperature=dto.temperature,
         )
 
-        return create_react_agent(
-            llm,
-            [magic_number_tool],
-            messages_modifier="""You are a helpful assistant.""",
-        )
+        return llm

tests/test_functional/fastapi_chat_completion_openai/test_server_openai.py

@@ -17,7 +17,7 @@ def openai_client():
 
 def test_chat_completion_invoke(openai_client):
     chat_completion = openai_client.chat.completions.create(
-        model="gpt-3.5-turbo",
+        model="gpt-4o-mini",
         messages=[
             {
                 "role": "user",
@@ -30,7 +30,7 @@ def test_chat_completion_invoke(openai_client):
 
 def test_chat_completion_stream(openai_client):
     chunks = openai_client.chat.completions.create(
-        model="gpt-3.5-turbo",
+        model="gpt-4o-mini",
         messages=[{"role": "user", "content": 'Say "This is a test"'}],
         stream=True,
     )
@@ -42,17 +42,3 @@ def test_chat_completion_stream(openai_client):
     stream_output = "".join(every_content)
 
     assert "This is a test" in stream_output
-
-
-def test_tool(openai_client):
-
-    chat_completion = openai_client.chat.completions.create(
-        model="gpt-3.5-turbo",
-        messages=[
-            {
-                "role": "user",
-                "content": 'Say "Magic number of 2"',
-            }
-        ],
-    )
-    assert "4" in chat_completion.choices[0].message.content

tests/test_functional/fastapi_chat_completion_react_agent_openai/my_openai_react_agent_factory.py

@@ -0,0 +1,30 @@
+from langchain_openai_api_bridge.core.base_agent_factory import BaseAgentFactory
+from langchain_core.runnables import Runnable
+from langchain_core.tools import tool
+from langgraph.prebuilt import create_react_agent
+from langchain_openai import ChatOpenAI
+
+from langchain_openai_api_bridge.core.create_agent_dto import CreateAgentDto
+
+
+@tool
+def magic_number_tool(input: int) -> int:
+    """Applies a magic function to an input."""
+    return input + 2
+
+
+class MyOpenAIReactAgentFactory(BaseAgentFactory):
+
+    def create_agent(self, dto: CreateAgentDto) -> Runnable:
+        llm = ChatOpenAI(
+            model=dto.model,
+            api_key=dto.api_key,
+            streaming=True,
+            temperature=dto.temperature,
+        )
+
+        return create_react_agent(
+            llm,
+            [magic_number_tool],
+            messages_modifier="""You are a helpful assistant.""",
+        )

tests/test_functional/fastapi_chat_completion_react_agent_openai/react_agent_server_openai.py

@@ -0,0 +1,37 @@
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from dotenv import load_dotenv, find_dotenv
+import uvicorn
+
+from langchain_openai_api_bridge.fastapi.langchain_openai_api_bridge_fastapi import (
+    LangchainOpenaiApiBridgeFastAPI,
+)
+from tests.test_functional.fastapi_chat_completion_react_agent_openai.my_openai_react_agent_factory import (
+    MyOpenAIReactAgentFactory,
+)
+
+_ = load_dotenv(find_dotenv())
+
+
+app = FastAPI(
+    title="Langchain Agent OpenAI API Bridge",
+    version="1.0",
+    description="OpenAI API exposing langchain agent",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+    expose_headers=["*"],
+)
+
+bridge = LangchainOpenaiApiBridgeFastAPI(
+    app=app, agent_factory_provider=lambda: MyOpenAIReactAgentFactory()
+)
+bridge.bind_openai_chat_completion()
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="localhost")

tests/test_functional/fastapi_chat_completion_react_agent_openai/test_react_agent_server_openai.py

@@ -0,0 +1,58 @@
+import pytest
+from openai import OpenAI
+from fastapi.testclient import TestClient
+from react_agent_server_openai import app
+
+
+test_api = TestClient(app)
+
+
+@pytest.fixture
+def openai_client():
+    return OpenAI(
+        base_url="http://testserver/openai/v1",
+        http_client=test_api,
+    )
+
+
+def test_chat_completion_invoke(openai_client):
+    chat_completion = openai_client.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[
+            {
+                "role": "user",
+                "content": 'Say "This is a test"',
+            }
+        ],
+    )
+    assert "This is a test" in chat_completion.choices[0].message.content
+
+
+def test_chat_completion_stream(openai_client):
+    chunks = openai_client.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": 'Say "This is a test"'}],
+        stream=True,
+    )
+    every_content = []
+    for chunk in chunks:
+        if chunk.choices and isinstance(chunk.choices[0].delta.content, str):
+            every_content.append(chunk.choices[0].delta.content)
+
+    stream_output = "".join(every_content)
+
+    assert "This is a test" in stream_output
+
+
+def test_tool(openai_client):
+
+    chat_completion = openai_client.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[
+            {
+                "role": "user",
+                "content": 'Say "Magic number of 2"',
+            }
+        ],
+    )
+    assert "4" in chat_completion.choices[0].message.content