From 0950fe76f62f5de4b758ed25c0fd900a5cc1b932 Mon Sep 17 00:00:00 2001
From: jemeza-codegen <jmeza@codegen.com>
Date: Wed, 12 Mar 2025 14:25:40 -0700
Subject: [PATCH] fix: increased max token limit

---
 .../swebench_agent_run/local_run.ipynb        | 78 ++++++++++++++++++-
 pyproject.toml                                |  3 +-
 src/codegen/extensions/langchain/agent.py     |  2 +-
 3 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/codegen-examples/examples/swebench_agent_run/local_run.ipynb b/codegen-examples/examples/swebench_agent_run/local_run.ipynb
index 0b212fa40..54d845c98 100644
--- a/codegen-examples/examples/swebench_agent_run/local_run.ipynb
+++ b/codegen-examples/examples/swebench_agent_run/local_run.ipynb
@@ -34,6 +34,82 @@
    "source": [
     "await run_eval(use_existing_preds=None, dataset=\"lite\", length=20, repo=\"django/django\", num_workers=10, model=\"claude-3-7-sonnet-latest\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from codegen.sdk.core.codebase import Codebase  # noqa: E402\n",
+    "\n",
+    "codebase = Codebase.from_repo(repo_full_name=\"codegen-sh/cloud\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# prompt = \"\"\"\n",
+    "# Please write integration tests for github tools. Below are paths to relevant files:\n",
+    "# Existing integration tests for linear (use these as an example):\n",
+    "#  - codegen-backend/codegen_integration/test_linear_tools.py\n",
+    "# Github Langchain Interface:\n",
+    "#  - codegen-backend/app/modal_app/gen/github_utils/tools.py\n",
+    "# Directory Containing github tools:\n",
+    "#  - codegen-backend/app/modal_app/gen/tools/github\n",
+    "# \"\"\"\n",
+    "prompt = \"\"\"\n",
+    "Add a notion integration to the codebase. View similar integrations for linear.\n",
+    "Please view the existing integrations for linear.\n",
+    "- src/codegen/extensions/clients/linear.py\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from codegen.agents.code_agent import CodeAgent\n",
+    "\n",
+    "agent = CodeAgent(codebase=codebase, tags=[\"local_test\"], model_name=\"claude-3-5-sonnet-latest\", model_provider=\"anthropic\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent.run(prompt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "codebase.reset()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -52,7 +128,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.0"
+   "version": "3.13.1"
   }
  },
  "nbformat": 4,
diff --git a/pyproject.toml b/pyproject.toml
index 441fe4954..2b13606e4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ readme = "README.md"
 # renovate: datasource=python-version depName=python
 requires-python = ">=3.12, <3.14"
 dependencies = [
-  "openai==1.66.1",
+  "openai==1.66.3",
   "tiktoken<1.0.0,>=0.5.1",
   "tabulate>=0.9.0,<1.0.0",
   "codeowners<1.0.0,>=0.6.0",
@@ -80,6 +80,7 @@ dependencies = [
   "colorlog>=6.9.0",
   "langsmith",
   "langchain-xai>=0.2.1",
+  "jwt>=1.3.1",
 ]
 
 license = { text = "Apache-2.0" }
diff --git a/src/codegen/extensions/langchain/agent.py b/src/codegen/extensions/langchain/agent.py
index fe44594b1..c1c796b9a 100644
--- a/src/codegen/extensions/langchain/agent.py
+++ b/src/codegen/extensions/langchain/agent.py
@@ -57,7 +57,7 @@ def create_codebase_agent(
     Returns:
         Initialized agent with message history
     """
-    llm = LLM(model_provider=model_provider, model_name=model_name, **kwargs)
+    llm = LLM(model_provider=model_provider, model_name=model_name, max_tokens=8192, **kwargs)
 
     # Get all codebase tools
     tools = [