refact: remove create_llm func

Author: samuelint

langchain_openai_api_bridge/core/base_agent_factory.py

@@ -1,15 +1,10 @@
 from abc import ABC, abstractmethod
 from langchain_core.runnables import Runnable
-from langchain_core.language_models import BaseChatModel
 from langchain_openai_api_bridge.core.create_agent_dto import CreateAgentDto
 
 
 class BaseAgentFactory(ABC):
 
     @abstractmethod
-    def create_agent(self, llm: BaseChatModel, dto: CreateAgentDto) -> Runnable:
-        pass
-
-    @abstractmethod
-    def create_llm(self, dto: CreateAgentDto) -> BaseChatModel:
+    def create_agent(self, dto: CreateAgentDto) -> Runnable:
         pass

langchain_openai_api_bridge/fastapi/chat_completion_router.py

@@ -30,8 +30,8 @@ async def assistant_retreive_thread_messages(
             api_key=api_key,
             temperature=request.temperature,
         )
-        llm = agent_factory.create_llm(dto=create_agent_dto)
-        agent = agent_factory.create_agent(llm=llm, dto=create_agent_dto)
+
+        agent = agent_factory.create_agent(dto=create_agent_dto)
 
         adapter = ChatCompletionCompatibleAPI.from_agent(agent, create_agent_dto.model)
 

langchain_openai_api_bridge/fastapi/internal_agent_factory.py

@@ -18,5 +18,4 @@ def create_agent(self, thread_run_dto: ThreadRunsDto, api_key: str) -> Runnable:
             temperature=thread_run_dto.temperature,
             assistant_id=thread_run_dto.assistant_id,
         )
-        llm = self.agent_factory.create_llm(dto=create_agent_dto)
-        return self.agent_factory.create_agent(llm=llm, dto=create_agent_dto)
+        return self.agent_factory.create_agent(dto=create_agent_dto)

tests/test_functional/fastapi_assistant_agent_anthropic/my_anthropic_agent_factory.py

@@ -4,7 +4,6 @@
 
 from langchain_openai_api_bridge.core.base_agent_factory import BaseAgentFactory
 from langchain_core.runnables import Runnable
-from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import tool
 from langgraph.prebuilt import create_react_agent
 
@@ -19,16 +18,15 @@ def magic_number_tool(input: int) -> int:
 
 class MyAnthropicAgentFactory(BaseAgentFactory):
 
-    def create_agent(self, llm: BaseChatModel, dto: CreateAgentDto) -> Runnable:
+    def create_agent(self, dto: CreateAgentDto) -> Runnable:
+        llm = AnthropicOpenAICompatibleChatModel(
+            model=dto.model,
+            max_tokens=1024,
+            streaming=True,
+        )
+
         return create_react_agent(
             llm,
             [magic_number_tool],
             messages_modifier="""You are a helpful assistant.""",
         )
-
-    def create_llm(self, dto: CreateAgentDto) -> Runnable:
-        return AnthropicOpenAICompatibleChatModel(
-            model=dto.model,
-            max_tokens=1024,
-            streaming=True,
-        )

tests/test_functional/fastapi_assistant_agent_custom_service/my_agent_factory.py

@@ -1,6 +1,5 @@
 from langchain_openai_api_bridge.core.base_agent_factory import BaseAgentFactory
 from langchain_core.runnables import Runnable
-from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import tool
 from langgraph.prebuilt import create_react_agent
 from langchain_openai import ChatOpenAI
@@ -16,17 +15,16 @@ def magic_number_tool(input: int) -> int:
 
 class MyAgentFactory(BaseAgentFactory):
 
-    def create_agent(self, llm: BaseChatModel, dto: CreateAgentDto) -> Runnable:
-        return create_react_agent(
-            llm,
-            [magic_number_tool],
-            messages_modifier="""You are a helpful assistant.""",
-        )
-
-    def create_llm(self, dto: CreateAgentDto) -> Runnable:
-        return ChatOpenAI(
+    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_assistant_agent_groq/my_groq_agent_factory.py

@@ -1,25 +1,21 @@
 from langchain_groq import ChatGroq
 from langchain_openai_api_bridge.core.base_agent_factory import BaseAgentFactory
 from langchain_core.runnables import Runnable
-from langchain_core.language_models import BaseChatModel
 from langgraph.prebuilt import create_react_agent
 
 from langchain_openai_api_bridge.core.create_agent_dto import CreateAgentDto
 
 
 class MyGroqAgentFactory(BaseAgentFactory):
 
-    def create_agent(self, llm: BaseChatModel, dto: CreateAgentDto) -> Runnable:
+    def create_agent(self, dto: CreateAgentDto) -> Runnable:
+        llm = ChatGroq(
+            model=dto.model,
+            streaming=False,  # Must be set to false. Groq does not support tool calling with stream at the moment (23/07/2024)
+        )
+
         return create_react_agent(
             llm,
             [],
             messages_modifier="""You are a helpful assistant.""",
         )
-
-    def create_llm(self, dto: CreateAgentDto) -> Runnable:
-        chat_model = ChatGroq(
-            model=dto.model,
-            streaming=False,  # Must be set to false. Groq does not support tool calling with stream at the moment (23/07/2024)
-        )
-
-        return chat_model

tests/test_functional/fastapi_assistant_agent_llamacpp/my_llamacpp_agent_factory.py

@@ -22,7 +22,9 @@ def magic_number_tool(input: int) -> int:
 
 class MyLlamacppAgentFactory(BaseAgentFactory):
 
-    def create_agent(self, llm: BaseChatModel, dto: CreateAgentDto) -> Runnable:
+    def create_agent(self, dto: CreateAgentDto) -> Runnable:
+        llm = self.create_llm(dto=dto)
+
         return create_react_agent(
             llm,
             [magic_number_tool],
@@ -45,6 +47,6 @@ def create_llm(self, dto: CreateAgentDto) -> BaseChatModel:
 
         return LLamacppOpenAICompatibleChatModel(
             llama=llama,
-            temperature=0,
+            temperature=dto.temperature or 0,
             streaming=True,
         )

tests/test_functional/fastapi_assistant_agent_openai/my_agent_factory.py

@@ -1,6 +1,5 @@
 from langchain_openai_api_bridge.core.base_agent_factory import BaseAgentFactory
 from langchain_core.runnables import Runnable
-from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import tool
 from langgraph.prebuilt import create_react_agent
 from langchain_openai import ChatOpenAI
@@ -16,17 +15,16 @@ def magic_number_tool(input: int) -> int:
 
 class MyAgentFactory(BaseAgentFactory):
 
-    def create_agent(self, llm: BaseChatModel, dto: CreateAgentDto) -> Runnable:
-        return create_react_agent(
-            llm,
-            [magic_number_tool],
-            messages_modifier="""You are a helpful assistant.""",
-        )
-
-    def create_llm(self, dto: CreateAgentDto) -> Runnable:
-        return ChatOpenAI(
+    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_anthropic/my_anthropic_agent_factory.py

@@ -1,6 +1,5 @@
 from langchain_openai_api_bridge.core.base_agent_factory import BaseAgentFactory
 from langchain_core.runnables import Runnable
-from langchain_core.language_models import BaseChatModel
 from langgraph.prebuilt import create_react_agent
 from langchain_anthropic import ChatAnthropic
 
@@ -9,15 +8,14 @@
 
 class MyAnthropicAgentFactory(BaseAgentFactory):
 
-    def create_agent(self, llm: BaseChatModel, dto: CreateAgentDto) -> Runnable:
+    def create_agent(self, dto: CreateAgentDto) -> Runnable:
+        llm = ChatAnthropic(
+            model=dto.model,
+            streaming=True,
+        )
+
         return create_react_agent(
             llm,
             [],
             messages_modifier="""You are a helpful assistant.""",
         )
-
-    def create_llm(self, dto: CreateAgentDto) -> Runnable:
-        return ChatAnthropic(
-            model=dto.model,
-            streaming=True,
-        )

tests/test_functional/fastapi_chat_completion_openai/my_openai_agent_factory.py

@@ -1,6 +1,5 @@
 from langchain_openai_api_bridge.core.base_agent_factory import BaseAgentFactory
 from langchain_core.runnables import Runnable
-from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import tool
 from langgraph.prebuilt import create_react_agent
 from langchain_openai import ChatOpenAI
@@ -16,17 +15,16 @@ def magic_number_tool(input: int) -> int:
 
 class MyOpenAIAgentFactory(BaseAgentFactory):
 
-    def create_agent(self, llm: BaseChatModel, dto: CreateAgentDto) -> Runnable:
-        return create_react_agent(
-            llm,
-            [magic_number_tool],
-            messages_modifier="""You are a helpful assistant.""",
-        )
-
-    def create_llm(self, dto: CreateAgentDto) -> Runnable:
-        return ChatOpenAI(
+    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.""",
+        )