python coding agent

langchain-streamlit-demo/app.py

@@ -28,6 +28,8 @@ from llm_resources import (
     get_texts_and_multiretriever,
 )
 from research_assistant.chain import get_chain as get_research_assistant_chain
+from python_coder import get_agent as get_python_agent
+
 
 __version__ = "2.0.1"
 
@@ -469,7 +471,19 @@ if st.session_state.llm:
                     "for quick facts, use duckduckgo instead.",
                 )
 
-                TOOLS = [research_assistant_tool] + default_tools
+                python_coder_agent = get_python_agent(st.session_state.llm)
+
+                python_coder_tool = Tool.from_function(
+                    func=lambda s: python_coder_agent.invoke(
+                        {"input": s},
+                        # config=get_config(callbacks),
+                    ),
+                    name="python-coder-assistant",
+                    description="this assistant writes Python code. give it clear instructions and requirements.",
+                )
+
+                TOOLS = [research_assistant_tool, python_coder_tool] + default_tools
+
                 if use_document_chat:
                     st.session_state.doc_chain = get_runnable(
                         use_document_chat,

langchain-streamlit-demo/python_coder.py

@@ -0,0 +1,208 @@
+"""langchain python coder-- requires black, ruff, and mypy."""
+
+import argparse
+import os
+import re
+import subprocess  # nosec
+import tempfile
+from importlib.util import find_spec
+
+from langchain.agents import initialize_agent, AgentType
+from langchain.agents.tools import Tool
+from langchain.chat_models import ChatOpenAI
+from langchain.llms.base import BaseLLM
+from langchain.prompts import ChatPromptTemplate
+from langchain.pydantic_v1 import BaseModel, validator, Field, ValidationError
+
+
+def strip_python_markdown_tags(text: str) -> str:
+    pat = re.compile(r"```python\n(.*)```", re.DOTALL)
+    code = pat.match(text)
+    if code:
+        return code.group(1)
+    else:
+        return text
+
+
+def format_black(filepath: str):
+    """Format a file with black."""
+    subprocess.run(  # nosec
+        f"black {filepath}",
+        stderr=subprocess.STDOUT,
+        text=True,
+        shell=True,
+        timeout=3,
+        check=False,
+    )
+
+
+def format_ruff(filepath: str):
+    """Run ruff format on a file."""
+    subprocess.run(  # nosec
+        f"ruff check --fix {filepath}",
+        shell=True,
+        text=True,
+        timeout=3,
+        universal_newlines=True,
+        check=False,
+    )
+
+    subprocess.run(  # nosec
+        f"ruff format {filepath}",
+        stderr=subprocess.STDOUT,
+        shell=True,
+        timeout=3,
+        text=True,
+        check=False,
+    )
+
+
+def check_ruff(filepath: str):
+    """Run ruff check on a file."""
+    subprocess.check_output(  # nosec
+        f"ruff check {filepath}",
+        stderr=subprocess.STDOUT,
+        shell=True,
+        timeout=3,
+        text=True,
+    )
+
+
+def check_mypy(filepath: str, strict: bool = True, follow_imports: str = "skip"):
+    """Run mypy on a file."""
+    cmd = f"mypy {'--strict' if strict else ''} --follow-imports={follow_imports} {filepath}"
+
+    subprocess.check_output(  # nosec
+        cmd,
+        stderr=subprocess.STDOUT,
+        shell=True,
+        text=True,
+        timeout=3,
+    )
+
+
+class PythonCode(BaseModel):
+    code: str = Field(
+        description="Python code conforming to ruff, black, and *strict* mypy standards.",
+    )
+
+    @validator("code")
+    @classmethod
+    def check_code(cls, v: str) -> str:
+        v = strip_python_markdown_tags(v).strip()
+        try:
+            with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
+                temp_file.write(v)
+                temp_file_path = temp_file.name
+
+            try:
+                # format with black and ruff
+                format_black(temp_file_path)
+                format_ruff(temp_file_path)
+            except subprocess.CalledProcessError:
+                pass
+
+            # update `v` with formatted code
+            with open(temp_file_path, "r") as temp_file:
+                v = temp_file.read()
+
+            # check
+            complaints = dict(ruff=None, mypy=None)
+
+            try:
+                check_ruff(temp_file_path)
+            except subprocess.CalledProcessError as e:
+                complaints["ruff"] = e.output
+
+            try:
+                check_mypy(temp_file_path)
+            except subprocess.CalledProcessError as e:
+                complaints["mypy"] = e.output
+
+            # raise ValueError if ruff or mypy had complaints
+            if any(complaints.values()):
+                code_str = f"```{temp_file_path}\n{v}```"
+                error_messages = [
+                    f"```{key}\n{value}```"
+                    for key, value in complaints.items()
+                    if value
+                ]
+                raise ValueError("\n\n".join([code_str] + error_messages))
+
+        finally:
+            os.remove(temp_file_path)
+        return v
+
+
+def check_code(code: str) -> str:
+    try:
+        code_obj = PythonCode(code=code)
+        return f"# LGTM\n# use the `submit` tool to submit this code:\n\n```python\n{code_obj.code}\n```"
+    except ValidationError as e:
+        return e.errors()[0]["msg"]
+
+
+prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            "You are a world class Python coder who uses black, ruff, and *strict* mypy for all of your code. "
+            "Provide complete, end-to-end Python code to meet the user's description/requirements. "
+            "Always `check` your code. When you're done, you must ALWAYS use the `submit` tool.",
+        ),
+        (
+            "human",
+            ": {input}",
+        ),
+    ],
+)
+
+check_code_tool = Tool.from_function(
+    check_code,
+    name="check-code",
+    description="Always check your code before submitting it!",
+)
+
+submit_code_tool = Tool.from_function(
+    lambda s: strip_python_markdown_tags(s),
+    name="submit-code",
+    description="THIS TOOL is the most important. use it to submit your code to the user who requested it... but be sure to `check` it first!",
+    return_direct=True,
+)
+
+tools = [check_code_tool, submit_code_tool]
+
+
+def get_agent(llm: BaseLLM, agent_type: AgentType = AgentType.OPENAI_FUNCTIONS):
+    return initialize_agent(
+        tools,
+        llm,
+        agent=agent_type,
+        verbose=True,
+        handle_parsing_errors=True,
+        prompt=prompt,
+        # return_intermediate_steps=True,
+    ) | (lambda output: output["output"])
+
+
+if __name__ == "__main__":
+    for lib in ("black", "ruff", "mypy"):
+        if find_spec(lib) is None:
+            raise ImportError(
+                "You must install black, ruff, and mypy to use this tool. "
+                "You can install them with `pip install black ruff mypy`.",
+            )
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", "-m", default="gpt-4-1106-preview")
+    parser.add_argument("instruction")
+    args = parser.parse_args()
+
+    instruction = args.instruction
+    model = args.model
+
+    llm = ChatOpenAI(model_name=model, temperature=0.0)
+    agent = get_agent(llm)
+
+    output = agent.invoke({"input": instruction}, config=dict(verbose=True))
+    print(output)

requirements.txt

@@ -1,20 +1,23 @@
 anthropic==0.7.7
 beautifulsoup4==4.12.2
+black==23.12.0
+certifi>=2023.7.22 # not directly required, pinned by Snyk to avoid a vulnerability
 duckduckgo-search==4.0.0
 faiss-cpu==1.7.4
 langchain==0.0.348
 langsmith==0.0.69
+mypy==1.7.1
 numexpr==2.8.8
 numpy>=1.22.2 # not directly required, pinned by Snyk to avoid a vulnerability
 openai==1.3.8
 pillow>=10.0.1 # not directly required, pinned by Snyk to avoid a vulnerability
 pyarrow>=14.0.1 # not directly required, pinned by Snyk to avoid a vulnerability
 pypdf==3.17.2
+requests>=2.31.0 # not directly required, pinned by Snyk to avoid a vulnerability
+ruff==0.1.8
 streamlit==1.29.0
 streamlit-feedback==0.1.3
 tiktoken==0.5.2
 tornado>=6.3.3 # not directly required, pinned by Snyk to avoid a vulnerability
 validators>=0.21.0 # not directly required, pinned by Snyk to avoid a vulnerability
 wikipedia==1.4.0
-certifi>=2023.7.22 # not directly required, pinned by Snyk to avoid a vulnerability
-requests>=2.31.0 # not directly required, pinned by Snyk to avoid a vulnerability