fix: remove MultiAIProvider (#555)

Author: jayhack

pyproject.toml

@@ -6,7 +6,6 @@ readme = "README.md"
 # renovate: datasource=python-version depName=python
 requires-python = ">=3.12, <3.14"
 dependencies = [
-  "backoff==2.2.1",
   "openai==1.63.2",
   "tiktoken<1.0.0,>=0.5.1",
   "tabulate>=0.9.0,<1.0.0",
@@ -31,14 +30,12 @@ dependencies = [
   "plotly>=5.24.0,<7.0.0",
   "humanize<5.0.0,>=4.10.0",
   "pytest-snapshot>=0.9.0",
-  "anthropic==0.23.1",
   "pyjson5==1.6.8",
   "mini-racer>=0.12.4",
   "rustworkx>=0.15.1",
   "typing-extensions>=4.12.2",
   "termcolor>=2.4.0",
   "sentry-sdk==2.22.0",
-  "tenacity>=9.0.0",
   "click>=8.1.7",
   "requests>=2.32.3",
   "lazy-object-proxy>=0.0.0",
@@ -72,6 +69,7 @@ dependencies = [
   "neo4j",
   "modal>=0.73.45",
   "slack-sdk",
+  "langchain-anthropic>=0.3.7",
 ]
 
 license = { text = "Apache-2.0" }

src/codegen/extensions/events/linear.py

@@ -4,7 +4,7 @@
 from typing import Callable
 
 import modal  # deptry: ignore
-from anthropic import BaseModel
+from pydantic import BaseModel
 
 from codegen.extensions.clients.linear import LinearClient
 from codegen.extensions.events.interface import EventHandlerManagerProtocol

src/codegen/extensions/langchain/agent.py

@@ -1,9 +1,10 @@
 """Demo implementation of an agent with Codegen tools."""
 
-from langchain.agents import AgentExecutor
+from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
 from langchain.hub import pull
 from langchain.tools import BaseTool
+from langchain_anthropic import ChatAnthropic
 from langchain_core.chat_history import InMemoryChatMessageHistory
 from langchain_core.messages import BaseMessage
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
@@ -29,7 +30,7 @@
 
 def create_codebase_agent(
     codebase: Codebase,
-    model_name: str = "gpt-4o",
+    model_name: str = "claude-3-5-sonnet-latest",
     temperature: float = 0,
     verbose: bool = True,
     chat_history: list[BaseMessage] = [],
@@ -46,8 +47,13 @@ def create_codebase_agent(
         Initialized agent with message history
     """
     # Initialize language model
-    llm = ChatOpenAI(
-        model_name=model_name,
+    # llm = ChatOpenAI(
+    #     model_name=model_name,
+    #     temperature=temperature,
+    # )
+
+    llm = ChatAnthropic(
+        model="claude-3-5-sonnet-latest",
         temperature=temperature,
     )
 
@@ -64,7 +70,8 @@ def create_codebase_agent(
         RevealSymbolTool(codebase),
         SemanticEditTool(codebase),
         SemanticSearchTool(codebase),
-        # CommitTool(codebase),
+        # =====[ Github Integration ]=====
+        # Enable Github integration
         # GithubCreatePRTool(codebase),
         # GithubViewPRTool(codebase),
         # GithubCreatePRCommentTool(codebase),
@@ -128,7 +135,12 @@ def create_codebase_agent(
     )
 
     # Create the agent
-    agent = OpenAIFunctionsAgent(
+    # agent = OpenAIFunctionsAgent(
+    #     llm=llm,
+    #     tools=tools,
+    #     prompt=prompt,
+    # )
+    agent = create_tool_calling_agent(
         llm=llm,
         tools=tools,
         prompt=prompt,

src/codegen/extensions/tools/reveal_symbol.py

@@ -3,17 +3,12 @@
 import tiktoken
 
 from codegen import Codebase
+from codegen.sdk.ai.utils import count_tokens
 from codegen.sdk.core.external_module import ExternalModule
 from codegen.sdk.core.import_resolution import Import
 from codegen.sdk.core.symbol import Symbol
 
 
-def count_tokens(text: str) -> int:
-    """Count the number of tokens in a string using GPT tokenizer."""
-    enc = tiktoken.get_encoding("cl100k_base")  # GPT-4 encoding
-    return len(enc.encode(text))
-
-
 def truncate_source(source: str, max_tokens: int) -> str:
     """Truncate source code to fit within max_tokens while preserving meaning.
 

src/codegen/extensions/tools/semantic_edit.py

@@ -3,12 +3,13 @@
 import difflib
 import re
 
+from langchain_anthropic import ChatAnthropic
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_openai import ChatOpenAI
 
 from codegen import Codebase
 
 from .tool_prompts import _HUMAN_PROMPT_DRAFT_EDITOR, _SYSTEM_PROMPT_DRAFT_EDITOR
+from .view_file import add_line_numbers
 
 
 def generate_diff(original: str, modified: str) -> str:
@@ -117,15 +118,25 @@ def semantic_edit(codebase: Codebase, filepath: str, edit_content: str, start: i
     original_content = file.content
     original_lines = original_content.split("\n")
 
+    # Check if file is too large for full edit
+    MAX_LINES = 300
+    if len(original_lines) > MAX_LINES and start == 1 and end == -1:
+        return {
+            "error": (
+                f"File is {len(original_lines)} lines long. For files longer than {MAX_LINES} lines, "
+                "please specify a line range using start and end parameters. "
+                "You may need to make multiple targeted edits."
+            ),
+            "status": "error",
+            "line_count": len(original_lines),
+        }
+
     # Handle append mode
     if start == -1 and end == -1:
         try:
             file.add_symbol_from_source(edit_content)
             codebase.commit()
 
-            # Analyze changes for append
-            new_lines = file.content.split("\n")
-
             return {"filepath": filepath, "content": file.content, "diff": generate_diff(original_content, file.content), "status": "success"}
         except Exception as e:
             msg = f"Failed to append content: {e!s}"
@@ -144,10 +155,10 @@ def semantic_edit(codebase: Codebase, filepath: str, edit_content: str, start: i
     system_message = _SYSTEM_PROMPT_DRAFT_EDITOR
     human_message = _HUMAN_PROMPT_DRAFT_EDITOR
     prompt = ChatPromptTemplate.from_messages([system_message, human_message])
-    llm = ChatOpenAI(
-        model="gpt-4o",
+    llm = ChatAnthropic(
+        model="claude-3-5-sonnet-latest",
         temperature=0,
-        max_tokens=10000,
+        max_tokens=5000,
     )
     chain = prompt | llm
     response = chain.invoke({"original_file_section": original_file_section, "edit_content": edit_content})
@@ -173,4 +184,4 @@ def semantic_edit(codebase: Codebase, filepath: str, edit_content: str, start: i
     file.edit(new_content)
     codebase.commit()
 
-    return {"filepath": filepath, "diff": diff, "status": "success"}
+    return {"filepath": filepath, "diff": diff, "status": "success", "new_content": add_line_numbers(new_content)}

src/codegen/extensions/tools/semantic_search.py

@@ -69,18 +69,12 @@ def semantic_search(
 
         # Format results with previews
         formatted_results = []
-        for filepath, score in results:
-            try:
-                file = codebase.get_file(filepath)
-                preview = file.content[:preview_length].replace("\n", " ").strip()
-                if len(file.content) > preview_length:
-                    preview += "..."
+        for file, score in results:
+            preview = file.content[:preview_length].replace("\n", " ").strip()
+            if len(file.content) > preview_length:
+                preview += "..."
 
-                formatted_results.append({"filepath": filepath, "score": float(score), "preview": preview})
-            except Exception as e:
-                # Skip files that can't be read
-                print(f"Warning: Could not read file {filepath}: {e}")
-                continue
+            formatted_results.append({"filepath": file.filepath, "score": float(score), "preview": preview})
 
         return {"status": "success", "query": query, "results": formatted_results}
 

src/codegen/gscli/generate/commands.py

@@ -11,7 +11,7 @@
 from codegen.gscli.generate.runner_imports import _generate_runner_imports
 from codegen.gscli.generate.system_prompt import get_system_prompt
 from codegen.gscli.generate.utils import LanguageType, generate_builtins_file
-from codegen.sdk.ai.helpers import AnthropicHelper
+from codegen.sdk.ai.client import get_openai_client
 from codegen.sdk.code_generation.changelog_generation import generate_changelog
 from codegen.sdk.code_generation.codegen_sdk_codebase import get_codegen_sdk_codebase
 from codegen.sdk.code_generation.doc_utils.generate_docs_json import generate_docs_json
@@ -201,9 +201,9 @@ def generate_codegen_sdk_docs(docs_dir: str) -> None:
 
 @generate.command()
 @click.option("--docs-dir", default="docs", required=False)
-@click.option("--anthropic-key", required=True)
+@click.option("--openai-key", required=True)
 @click.option("--complete", is_flag=True, help="Generate a complete changelog for the codegen_sdk API")
-def changelog(docs_dir: str, anthropic_key: str, complete: bool = False) -> None:
+def changelog(docs_dir: str, openai_key: str, complete: bool = False) -> None:
     """Generate the changelog for the codegen_sdk API and update the changelog.mdx file"""
     print(colored("Generating changelog", "green"))
     header = """---
@@ -212,8 +212,8 @@ def changelog(docs_dir: str, anthropic_key: str, complete: bool = False) -> None
 iconType: "solid"
 ---
 """
-    # Generate the changelog for the codegen_sdk API and update the changelog.mdx file
-    client = AnthropicHelper(anthropic_key=anthropic_key, cache=True, openai_anthropic_translation=False)
+
+    client = get_openai_client(openai_key)
 
     if complete:
         entire_release_history = generate_changelog(client)

src/codegen/runner/sandbox/runner.py

@@ -11,8 +11,8 @@
 from codegen.sdk.codebase.config import CodebaseConfig, ProjectConfig, SessionOptions
 from codegen.sdk.codebase.factory.codebase_factory import CodebaseType
 from codegen.sdk.core.codebase import Codebase
-from codegen.sdk.secrets import Secrets
 from codegen.shared.compilation.string_to_code import create_execute_function_from_codeblock
+from codegen.shared.configs.models.secrets import SecretsConfig
 from codegen.shared.configs.session_configs import config
 from codegen.shared.performance.stopwatch_utils import stopwatch
 
@@ -47,7 +47,7 @@ async def warmup(self) -> None:
     async def _build_graph(self) -> Codebase:
         logger.info("> Building graph...")
         projects = [ProjectConfig(programming_language=self.repo.language, repo_operator=self.op, base_path=self.repo.base_path, subdirectories=self.repo.subdirectories)]
-        secrets = Secrets(openai_key=config.secrets.openai_api_key)
+        secrets = SecretsConfig(openai_api_key=config.secrets.openai_api_key)
         codebase_config = CodebaseConfig(secrets=secrets, feature_flags=config.feature_flags.codebase)
         return Codebase(projects=projects, config=codebase_config)
 

src/codegen/sdk/ai/client.py

@@ -0,0 +1,5 @@
+from openai import OpenAI
+
+
+def get_openai_client(key: str) -> OpenAI:
+    return OpenAI(api_key=key)