Use init_chat_model() helper method to initialise models (#87)

* Upgrade langchain version to 0.2.8

* Use new init_chat_model function to initialise any models. Remove models.py since its redundant

* Add migration file to convert all provider names in member table to new name format

* Update default provider and model names when creating new member

* Update frontend available model provider names

* Update test_members to use new provider name

backend/app/alembic/versions/6e7c33ddf30f_rename_provider_names_in_members_table_.py

@@ -0,0 +1,40 @@
+"""Rename provider names in member table to new name
+
+Revision ID: 6e7c33ddf30f
+Revises: 0a354b5c6f6c
+Create Date: 2024-07-27 04:29:51.886906
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '6e7c33ddf30f'
+down_revision = '0a354b5c6f6c'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # Mapping of old provider names to new provider names
+    mapping = {
+        'ChatOpenAI': 'openai',
+        'ChatAnthropic': 'anthropic',
+    }
+
+    # Rename each provider name according to the mapping
+    for old_name, new_name in mapping.items():
+        op.execute(f"UPDATE member SET provider = '{new_name}' WHERE provider = '{old_name}'")
+
+
+def downgrade():
+    # Mapping of new provider names back to old provider names
+    mapping = {
+        'openai': 'ChatOpenAI',
+        'anthropic': 'ChatAnthropic',
+    }
+
+    # Revert each provider name according to the mapping
+    for new_name, old_name in mapping.items():
+        op.execute(f"UPDATE member SET provider = '{old_name}' WHERE provider = '{new_name}'")

backend/app/core/graph/members.py

@@ -1,6 +1,7 @@
 from collections.abc import Mapping, Sequence
 from typing import Annotated, Any
 
+from langchain.chat_models import init_chat_model
 from langchain.tools.retriever import create_retriever_tool
 from langchain_core.messages import AIMessage, AnyMessage
 from langchain_core.output_parsers.openai_tools import JsonOutputKeyToolsParser
@@ -16,7 +17,6 @@
 from pydantic import BaseModel, Field
 from typing_extensions import NotRequired, TypedDict
 
-from app.core.graph.models import all_models
 from app.core.graph.rag.qdrant import QdrantStore
 from app.core.graph.skills import managed_skills
 from app.core.graph.skills.api_tool import dynamic_api_tool
@@ -147,12 +147,12 @@ class ReturnTeamState(TypedDict):
 
 class BaseNode:
     def __init__(self, provider: str, model: str, temperature: float):
-        self.model = all_models[provider](
-            model=model, temperature=temperature, streaming=True
-        )  # type: ignore[call-arg]
-        self.final_answer_model = all_models[provider](
-            model=model, temperature=0, streaming=True
-        )  # type: ignore[call-arg]
+        self.model = init_chat_model(
+            model, model_provider=provider, temperature=temperature, streaming=True
+        )
+        self.final_answer_model = init_chat_model(
+            model, model_provider=provider, temperature=0, streaming=True
+        )
 
     def tag_with_name(self, ai_message: AIMessage, name: str) -> AIMessage:
         """Tag a name to the AI message"""

backend/app/models.py

@@ -251,8 +251,8 @@ class MemberBase(SQLModel):
     position_x: float
     position_y: float
     source: int | None = None
-    provider: str = "ChatOpenAI"
-    model: str = "gpt-3.5-turbo"
+    provider: str = "openai"
+    model: str = "gpt-4o-mini"
     temperature: float = 0.7
     interrupt: bool = False
 

backend/app/tests/api/routes/test_members.py

@@ -33,7 +33,7 @@ def test_read_members(
         "position_x": 0.0,
         "position_y": 0.0,
         "source": None,
-        "provider": "ChatOpenAI",
+        "provider": "openai",
         "model": "gpt-3.5-turbo",
         "temperature": 0.7,
         "interrupt": False,
@@ -70,7 +70,7 @@ def test_read_member(
         "position_x": 0.0,
         "position_y": 0.0,
         "source": None,
-        "provider": "ChatOpenAI",
+        "provider": "openai",
         "model": "gpt-3.5-turbo",
         "temperature": 0.7,
         "interrupt": False,
@@ -107,7 +107,7 @@ def test_create_member(
         "position_x": 0.0,
         "position_y": 0.0,
         "source": None,
-        "provider": "ChatOpenAI",
+        "provider": "openai",
         "model": "gpt-3.5-turbo",
         "temperature": 0.7,
         "interrupt": False,
@@ -138,7 +138,7 @@ def test_create_member_duplicate_name(
         "position_x": 0.0,
         "position_y": 0.0,
         "source": None,
-        "provider": "ChatOpenAI",
+        "provider": "openai",
         "model": "gpt-3.5-turbo",
         "temperature": 0.7,
         "interrupt": False,
@@ -160,7 +160,7 @@ def test_create_member_duplicate_name(
         "position_x": 0.0,
         "position_y": 0.0,
         "source": None,
-        "provider": "ChatOpenAI",
+        "provider": "openai",
         "model": "gpt-3.5-turbo",
         "temperature": 0.7,
         "interrupt": False,
@@ -187,7 +187,7 @@ def test_update_member(
         "position_x": 0.0,
         "position_y": 0.0,
         "source": None,
-        "provider": "ChatOpenAI",
+        "provider": "openai",
         "model": "gpt-3.5-turbo",
         "temperature": 0.7,
         "interrupt": False,
@@ -226,7 +226,7 @@ def test_delete_member(
         "position_x": 0.0,
         "position_y": 0.0,
         "source": None,
-        "provider": "ChatOpenAI",
+        "provider": "openai",
         "model": "gpt-3.5-turbo",
         "temperature": 0.7,
         "interrupt": False,

backend/pyproject.toml

@@ -29,7 +29,7 @@ langgraph = "0.1.9"
 langserve = {extras = ["server"], version = "^0.0.51"}
 langchain-openai = "0.1.17"
 grandalf = "^0.8"
-langchain = "0.2.7"
+langchain = "0.2.8"
 langchain-community = "0.2.7"
 duckduckgo-search = "6.1.0"
 wikipedia = "^1.4.0"

frontend/src/components/Members/EditMember.tsx

@@ -52,14 +52,12 @@ const customSelectOption = {
 
 // TODO: Place this somewhere else.
 const AVAILABLE_MODELS = {
-  ChatOpenAI: ["gpt-4o-mini", "gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo"],
-  ChatAnthropic: [
+  openai: ["gpt-4o-mini", "gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo"],
+  anthropic: [
     "claude-3-opus-20240229",
     "claude-3-sonnet-20240229",
     "claude-3-haiku-20240307",
   ],
-  // ChatCohere: ["command"],
-  // ChatGoogleGenerativeAI: ["gemini-pro"],
 }
 
 type ModelProvider = keyof typeof AVAILABLE_MODELS