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