Remove langchain!!

Author: scosman

app/desktop/studio_server/provider_api.py

@@ -4,6 +4,7 @@
 from datetime import datetime, timedelta
 from typing import Any, Dict, List
 
+import litellm
 import openai
 import requests
 from fastapi import FastAPI, HTTPException
@@ -24,7 +25,6 @@
 from kiln_ai.datamodel.registry import all_projects
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
-from langchain_aws import ChatBedrockConverse
 from pydantic import BaseModel, Field
 
 logger = logging.getLogger(__name__)
@@ -611,42 +611,34 @@ async def connect_bedrock(key_data: dict):
             detail="Access Key or Secret Key not found",
         )
     try:
-        # Langchain API is not good... need to use env vars. Pop these in finally block
-        os.environ["AWS_ACCESS_KEY_ID"] = access_key
-        os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
-
-        # Fake model, but will get a credential error before AWS realizes it's wrong
-        # Ensures we don't spend money on a test call
-        llm = ChatBedrockConverse(
-            model="fake_model",
-            region_name="us-west-2",
+        # Test credentials request, but invalid model so we don't use tokens
+        await litellm.acompletion(
+            model="bedrock/ai21.jamba-1-5-mini-v9999.8888",
+            aws_region_name="us-west-2",
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key,
+            messages=[{"role": "user", "content": "Hello, how are you?"}],
         )
-        llm.invoke("Hello, how are you?")
     except Exception as e:
-        # Check for specific error messages indicating invalid credentials
-        if "UnrecognizedClientException" in str(
-            e
-        ) or "InvalidSignatureException" in str(e):
+        # Improve error message if it's a confirmed authentication error
+        if isinstance(e, litellm.exceptions.AuthenticationError):
             return JSONResponse(
                 status_code=401,
                 content={
                     "message": "Failed to connect to Bedrock. Invalid credentials."
                 },
             )
-        else:
-            # We still expect an error (fake model), but not for invalid credentials which means success!
+        # If it's a bad request, it's a valid key (but the model is fake)
+        if isinstance(e, litellm.exceptions.BadRequestError):
             Config.shared().bedrock_access_key = access_key
             Config.shared().bedrock_secret_key = secret_key
             return JSONResponse(
                 status_code=200,
                 content={"message": "Connected to Bedrock"},
             )
+        # Unknown error, raise it
+        raise e
 
-    finally:
-        os.environ.pop("AWS_ACCESS_KEY_ID")
-        os.environ.pop("AWS_SECRET_ACCESS_KEY")
-
-    # we shouldn't get here, but if we do, something went wrong
     return JSONResponse(
         status_code=400,
         content={"message": "Unknown Bedrock Error"},

conftest.py

@@ -1,10 +1,16 @@
 from unittest.mock import patch
 
+import litellm
 import pytest
 from dotenv import load_dotenv
 from kiln_ai.utils.config import Config
 
 
+@pytest.fixture(autouse=True)
+def _clear_httpx_clients() -> None:
+    litellm.in_memory_llm_clients_cache.flush_cache()
+
+
 @pytest.fixture(scope="session", autouse=True)
 def load_env():
     load_dotenv()

libs/core/kiln_ai/adapters/adapter_registry.py

@@ -3,7 +3,6 @@
 from kiln_ai import datamodel
 from kiln_ai.adapters.ml_model_list import ModelProviderName
 from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig, BaseAdapter
-from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
 from kiln_ai.adapters.model_adapters.openai_model_adapter import (
     OpenAICompatibleAdapter,
     OpenAICompatibleConfig,
@@ -62,7 +61,6 @@ def adapter_for_task(
                 prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
             )
-        # Use LangchainAdapter for the rest
         case ModelProviderName.groq:
             return OpenAICompatibleAdapter(
                 kiln_task=kiln_task,
@@ -138,19 +136,6 @@ def adapter_for_task(
                     provider_name=provider,
                 ),
             )
-        case ModelProviderName.azure:
-            return OpenAICompatibleAdapter(
-                kiln_task=kiln_task,
-                prompt_id=prompt_id,
-                base_adapter_config=base_adapter_config,
-                config=OpenAICompatibleConfig(
-                    litellm_provider_name="azure_ai",
-                    api_key=Config.shared().azure_ai_api_key,
-                    # base_url=Config.shared().azure_ai_api_base,
-                    model_name=model_name,
-                    provider_name=provider,
-                ),
-            )
         case ModelProviderName.azure_openai:
             return OpenAICompatibleAdapter(
                 kiln_task=kiln_task,

libs/core/kiln_ai/adapters/ml_model_list.py

@@ -28,8 +28,6 @@ class ModelProviderName(str, Enum):
     openai_compatible = "openai_compatible"
     anthropic = "anthropic"
     gemini_api = "gemini_api"
-    # They are going to change the name 16 times so don't include "ai studio" or "foundry" in the name :P
-    azure = "azure"
     azure_openai = "azure_openai"
 
 
@@ -876,10 +874,6 @@ class KilnModel(BaseModel):
                 supports_data_gen=False,
                 provider_options={"model": "microsoft/phi-4"},
             ),
-            KilnModelProvider(
-                name=ModelProviderName.azure,
-                provider_options={"model": "Phi-4"},
-            ),
         ],
     ),
     # Gemma 2 2.6b

libs/core/kiln_ai/adapters/model_adapters/__init__.py

@@ -7,12 +7,10 @@
 
 from . import (
     base_adapter,
-    langchain_adapters,
     openai_model_adapter,
 )
 
 __all__ = [
     "base_adapter",
-    "langchain_adapters",
     "openai_model_adapter",
 ]