Build an Autonomous Coding Agent with GitHub CI/CD and LangGraph

# requirements.txt
# langchain-openai==0.1.7
# langgraph==0.2.0
# python-dotenv==1.0.1
# e2b==0.1.10
# PyGithub==2.2.0
# tenacity==8.2.3

# .env (template)
# OPENAI_API_KEY="your_openai_api_key_here"
# E2B_API_KEY="your_e2b_api_key_here"
# GITHUB_TOKEN="your_github_personal_access_token_here"
# GITHUB_REPO_OWNER="your_github_username_or_org"
# GITHUB_REPO_NAME="your_repo_name"

# Optional: LangSmith for observability
# LANGCHAIN_TRACING_V2="true"
# LANGCHAIN_API_KEY="your_langsmith_api_key_here"
# LANGCHAIN_PROJECT="Autonomous Coding Agent"

# main.py
import os
import logging
from dotenv import load_dotenv
from typing import TypedDict, List, Annotated
from langgraph.graph import StateGraph, END
from agent import create_agent_graph, AgentState

load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def run_agent(task: str) -> str:
    """
    Initializes and runs the autonomous coding agent with a given task.
    """
    openai_api_key = os.environ.get("OPENAI_API_KEY")
    e2b_api_key = os.environ.get("E2B_API_KEY")
    github_token = os.environ.get("GITHUB_TOKEN")
    github_repo_owner = os.environ.get("GITHUB_REPO_OWNER")
    github_repo_name = os.environ.get("GITHUB_REPO_NAME")

    if not all([openai_api_key, e2b_api_key, github_token, github_repo_owner, github_repo_name]):
        logger.error("Missing one or more required environment variables.")
        raise ValueError("Ensure OPENAI_API_KEY, E2B_API_KEY, GITHUB_TOKEN, GITHUB_REPO_OWNER, GITHUB_REPO_NAME are set.")

    workflow = create_agent_graph(openai_api_key)
    
    inputs: AgentState = {
        "task": task, 
        "chat_history": [],
        "plan": "",
        "code": "",
        "test_code": "",
        "test_results": "",
        "pull_request_url": "",
        "iterations": 0,
        "max_iterations": 5, # Default, can be overridden by plan_task
        "output": "",
        "pr_number": 0, # Placeholder
        "github_repo_owner": github_repo_owner,
        "github_repo_name": github_repo_name
    }
    
    logger.info(f"Starting agent for task: {task}")
    final_output = None
    try:
        for s in workflow.stream(inputs, {"configurable": {"thread_id": "1"}}):
            logger.info(f"Agent step: {s}")
            final_output = s
        
        if final_output:
            # final_output is {"node_name": state_dict} from the last streamed step
            last_state = next(iter(final_output.values()), {}) if isinstance(final_output, dict) else {}
            return last_state.get("output", str(final_output))
        else:
            return "Agent finished without explicit output."

    except Exception as e:
        logger.exception(f"An error occurred during agent execution: {e}")
        return f"Agent execution failed: {e}"

if __name__ == "__main__":
    test_task = "Create a Python function `multiply(a, b)` that returns the product of two numbers. Include a docstring and a simple unit test for it. The function should be in a file named `src/generated_code.py` and tests in `tests/test_generated_code.py`."
    print("--- Running autonomous coding agent ---")
    result = run_agent(test_task)
    print("
--- Agent Finished ---")
    print(f"Final Result: {result}")

# agent.py
from typing import TypedDict, List, Annotated
import operator
import logging
from langgraph.graph import StateGraph, END
from langchain_core.messages import BaseMessage

logger = logging.getLogger(__name__)

class AgentState(TypedDict):
    """
    Represents the state of our agent.

    Attributes:
        task: The original task description.
        plan: The current high-level plan for execution.
        code: The generated code.
        test_code: The generated test code.
        test_results: Results from running tests.
        pull_request_url: URL of the created pull request.
        chat_history: A list of messages representing the conversation history.
        iterations: Number of self-correction iterations.
        max_iterations: Maximum allowed self-correction iterations.
        output: The final output message from the agent.
        pr_number: The number of the created pull request (for monitoring).
        github_repo_owner: Owner of the GitHub repository.
        github_repo_name: Name of the GitHub repository.
    """
    task: str
    plan: str
    code: str
    test_code: str
    test_results: str
    pull_request_url: str
    chat_history: Annotated[List[BaseMessage], operator.add]
    iterations: int
    max_iterations: int
    output: str
    pr_number: int
    github_repo_owner: str
    github_repo_name: str

def create_agent_graph(openai_api_key: str):
    """
    Creates and compiles the LangGraph workflow for the autonomous coding agent.
    """
    workflow = StateGraph(AgentState)

    # Define nodes here (will be added in subsequent phases)
    # Define edges here (will be added in subsequent phases)

    logger.info("LangGraph workflow initialized.")
    return workflow.compile()

# tools.py
import os
import logging
from e2b import Sandbox
from tenacity import retry, stop_after_attempt, wait_exponential
from typing import Dict, Any

logger = logging.getLogger(__name__)

# Sandbox instance will be initialized once and reused
_sandbox_instance: Sandbox | None = None

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def get_e2b_sandbox() -> Sandbox:
    """
    Retrieves or initializes the E2B sandbox instance.
    Retries on connection errors.
    """
    global _sandbox_instance
    if _sandbox_instance is None:
        e2b_api_key = os.environ.get("E2B_API_KEY")
        if not e2b_api_key:
            logger.error("E2B_API_KEY environment variable not set.")
            raise ValueError("E2B_API_KEY environment variable not set.")
        
        try:
            _sandbox_instance = Sandbox(api_key=e2b_api_key)
            logger.info("E2B Sandbox initialized successfully.")
        except Exception as e:
            logger.error(f"Failed to initialize E2B Sandbox: {e}")
            raise
    return _sandbox_instance

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def execute_python_code(code: str) -> Dict[str, Any]:
    """
    Executes Python code in the E2B sandbox and returns the stdout, stderr, and error status.
    Retries on sandbox connection or execution errors.
    """
    sandbox = get_e2b_sandbox()
    try:
        _escaped = code.replace("'", "\'")
        proc = sandbox.process.start(
            cmd=f"python -c '{_escaped}'",  # Escape single quotes for shell
            timeout=10
        )
        proc.wait()
        
        stdout = proc.stdout
        stderr = proc.stderr
        exit_code = proc.exit_code

        logger.info(f"Code execution finished with exit code {exit_code}")
        logger.debug(f"Stdout: {stdout}")
        logger.debug(f"Stderr: {stderr}")

        if exit_code != 0:
            logger.warning(f"Code execution failed. Stderr: {stderr}")
            return {"status": "error", "stdout": stdout, "stderr": stderr, "exit_code": exit_code}
        else:
            return {"status": "success", "stdout": stdout, "stderr": stderr, "exit_code": exit_code}
    except Exception as e:
        logger.exception(f"Error executing code in sandbox: {e}")
        return {"status": "runtime_error", "stdout": "", "stderr": str(e), "exit_code": -1}

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def write_file_to_sandbox(file_path: str, content: str) -> Dict[str, Any]:
    """
    Writes content to a specified file path in the E2B sandbox.
    """
    sandbox = get_e2b_sandbox()
    try:
        sandbox.filesystem.write(file_path, content)
        logger.info(f"File '{file_path}' written to sandbox.")
        return {"status": "success", "message": f"File '{file_path}' written."}
    except Exception as e:
        logger.error(f"Failed to write file '{file_path}' to sandbox: {e}")
        return {"status": "error", "message": str(e)}

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def read_file_from_sandbox(file_path: str) -> Dict[str, Any]:
    """
    Reads content from a specified file path in the E2B sandbox.
    """
    sandbox = get_e2b_sandbox()
    try:
        content = sandbox.filesystem.read(file_path)
        logger.info(f"File '{file_path}' read from sandbox.")
        return {"status": "success", "content": content}
    except Exception as e:
        logger.error(f"Failed to read file '{file_path}' from sandbox: {e}")
        return {"status": "error", "message": str(e)}

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def list_files_in_sandbox(path: str = ".") -> Dict[str, Any]:
    """
    Lists files in a given directory in the E2B sandbox.
    """
    sandbox = get_e2b_sandbox()
    try:
        files = sandbox.filesystem.ls(path)
        file_names = [f.name for f in files]
        logger.info(f"Listed files in '{path}': {file_names}")
        return {"status": "success", "files": file_names}
    except Exception as e:
        logger.error(f"Failed to list files in '{path}' in sandbox: {e}")
        return {"status": "error", "message": str(e)}

# agent.py (updated)
from typing import TypedDict, List, Annotated, Callable
import operator
import logging
import json
import os
from langgraph.graph import StateGraph, END
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from tenacity import retry, stop_after_attempt, wait_exponential

from tools import execute_python_code, write_file_to_sandbox, read_file_from_sandbox, list_files_in_sandbox

logger = logging.getLogger(__name__)

class AgentState(TypedDict):
    """
    Represents the state of our agent.

    Attributes:
        task: The original task description.
        plan: The current high-level plan for execution.
        code: The generated code.
        test_code: The generated test code.
        test_results: Results from running tests.
        pull_request_url: URL of the created pull request.
        chat_history: A list of messages representing the conversation history.
        iterations: Number of self-correction iterations.
        max_iterations: Maximum allowed self-correction iterations.
        output: The final output message from the agent.
        pr_number: The number of the created pull request (for monitoring).
        github_repo_owner: Owner of the GitHub repository.
        github_repo_name: Name of the GitHub repository.
    """
    task: str
    plan: str
    code: str
    test_code: str
    test_results: str
    pull_request_url: str
    chat_history: Annotated[List[BaseMessage], operator.add]
    iterations: int
    max_iterations: int
    output: str
    pr_number: int
    github_repo_owner: str
    github_repo_name: str

class Plan(BaseModel):
    """Plan to write and test code."""
    steps: List[str] = Field(description="List of steps to complete the task.")

class CodeResponse(BaseModel):
    """Response containing generated code."""
    code: str = Field(description="The generated Python code.")

class TestResponse(BaseModel):
    """Response containing generated test code."""
    test_code: str = Field(description="The generated Python test code.")

class Reflection(BaseModel):
    """Reflection on test results and plan for next steps."""
    thought: str = Field(description="Detailed thought process on why tests failed or what needs improvement.")
    new_plan_steps: List[str] = Field(description="Revised plan steps based on reflection.")
    is_complete: bool = Field(description="True if the task is considered complete and ready for PR.")

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def _get_llm(openai_api_key: str):
    """Initializes and returns the OpenAI ChatModel."""
    return ChatOpenAI(model="gpt-4o", temperature=0.7, api_key=openai_api_key)

def _get_prompt(system_message: str):
    """Creates a chat prompt template."""
    return ChatPromptTemplate.from_messages([
        ("system", system_message),
        ("placeholder", "{chat_history}"),
        ("human", "{task}"),
        ("human", "{current_state}"), # Optional: pass current state for more context
    ])

def plan_task(state: AgentState, config: dict) -> AgentState:
    """Agent node to generate a plan for the coding task."""
    logger.info("--- Planning Task ---")
    llm = _get_llm(config["configurable"]["openai_api_key"])
    
    plan_prompt = _get_prompt(
        "You are an expert software engineer. Given a task, create a detailed plan with specific steps to implement the required code and unit tests. "
        "Break down complex problems into smaller, manageable steps. Focus on implementation details, testing strategy, and potential edge cases. "
        "Output your plan as a JSON object with a 'steps' key containing a list of strings."
    ).partial(current_state="")

    plan_chain = plan_prompt | llm.with_structured_output(Plan)
    
    try:
        plan_obj = plan_chain.invoke({
            "chat_history": state["chat_history"],
            "task": state["task"]
        })
        plan_str = "
".join(f"- {step}" for step in plan_obj.steps)
        logger.info(f"Generated Plan:
{plan_str}")

        github_repo_owner = os.environ.get("GITHUB_REPO_OWNER")
        github_repo_name = os.environ.get("GITHUB_REPO_NAME")
        if not github_repo_owner or not github_repo_name:
            logger.error("GITHUB_REPO_OWNER or GITHUB_REPO_NAME not set.")
            raise ValueError("GitHub repository details not configured.")

        return {
            "plan": plan_str,
            "chat_history": [HumanMessage(content=state["task"]), AIMessage(content=f"Plan: {plan_str}")],
            "iterations": 0,
            "max_iterations": 5,
            "github_repo_owner": github_repo_owner,
            "github_repo_name": github_repo_name
        }
    except Exception as e:
        logger.exception(f"Error in plan_task: {e}")
        return {"output": f"Failed to plan task: {e}"}

def generate_code(state: AgentState, config: dict) -> AgentState:
    """Agent node to generate code based on the plan."""
    logger.info("--- Generating Code ---")
    llm = _get_llm(config["configurable"]["openai_api_key"])
    
    code_prompt = _get_prompt(
        "You are an expert Python programmer. Given the task and current plan, write the complete Python code. "
        "Focus on writing clean, correct, and robust code. Do not include test code. "
        "Output ONLY the Python code, ensuring it's syntactically correct and ready to be executed. "
        "Wrap your code in a JSON object with a 'code' key."
    )
    code_chain = code_prompt | llm.with_structured_output(CodeResponse)

    try:
        code_response = code_chain.invoke({
            "chat_history": state["chat_history"],
            "task": state["task"],
            "current_state": f"Current Plan:
{state['plan']}"
        })
        logger.info(f"Generated Code:
{code_response.code}")
        return {
            "code": code_response.code,
            "chat_history": state["chat_history"] + [AIMessage(content=f"Generated Code:
```python
{code_response.code}
```")]
        }
    except Exception as e:
        logger.exception(f"Error in generate_code: {e}")
        return {"output": f"Failed to generate code: {e}"}

def generate_tests(state: AgentState, config: dict) -> AgentState:
    """Agent node to generate unit tests for the generated code."""
    logger.info("--- Generating Tests ---")
    llm = _get_llm(config["configurable"]["openai_api_key"])
    
    test_prompt = _get_prompt(
        "You are an expert Python developer focused on test-driven development. "
        "Given the task and the generated code, write comprehensive unit tests. "
        "Ensure tests cover functionality, edge cases, and error conditions. "
        "Use `unittest` or `pytest` if applicable, but for simple functions, direct assertions are fine. "
        "Output ONLY the Python test code. It should be executable and import the main code if necessary. "
        "Wrap your test code in a JSON object with a 'test_code' key."
    )
    test_chain = test_prompt | llm.with_structured_output(TestResponse)

    try:
        test_response = test_chain.invoke({
            "chat_history": state["chat_history"],
            "task": state["task"],
            "current_state": f"Current Plan:
{state['plan']}

Generated Code:
{state['code']}"
        })
        logger.info(f"Generated Test Code:
{test_response.test_code}")
        return {
            "test_code": test_response.test_code,
            "chat_history": state["chat_history"] + [AIMessage(content=f"Generated Tests:
```python
{test_response.test_code}
```")]
        }
    except Exception as e:
        logger.exception(f"Error in generate_tests: {e}")
        return {"output": f"Failed to generate tests: {e}"}

def execute_tests(state: AgentState) -> AgentState:
    """Agent node to execute generated code and tests in the E2B sandbox."""
    logger.info("--- Executing Code and Tests ---")
    
    write_code_result = write_file_to_sandbox("main.py", state["code"])
    if write_code_result["status"] == "error":
        return {"output": f"Failed to write main.py: {write_code_result['message']}"}

    write_test_result = write_file_to_sandbox("test_main.py", state["test_code"])
    if write_test_result["status"] == "error":
        return {"output": f"Failed to write test_main.py: {write_test_result['message']}"}

    test_execution_result = execute_python_code("import unittest; from test_main import *; unittest.main(exit=False)")
    
    if test_execution_result["status"] == "runtime_error":
        logger.error(f"Sandbox runtime error during test execution: {test_execution_result['stderr']}")
        return {
            "test_results": f"Sandbox runtime error: {test_execution_result['stderr']}",
            "chat_history": state["chat_history"] + [AIMessage(content=f"Test Execution Failed (Sandbox Error):
{test_execution_result['stderr']}")]
        }
    
    stdout = test_execution_result["stdout"]
    stderr = test_execution_result["stderr"]
    exit_code = test_execution_result["exit_code"]

    full_output = f"Stdout:
{stdout}

Stderr:
{stderr}

Exit Code: {exit_code}"
    logger.info(f"Test Execution Result:
{full_output}")

    return {
        "test_results": full_output,
        "chat_history": state["chat_history"] + [AIMessage(content=f"Test Execution Result:
```
{full_output}
```")]
    }

def reflect_on_results(state: AgentState, config: dict) -> AgentState:
    """Agent node to reflect on test results and potentially update the plan for self-correction."""
    logger.info("--- Reflecting on Results ---")
    llm = _get_llm(config["configurable"]["openai_api_key"])
    
    reflection_prompt = _get_prompt(
        "You are an expert debugger and code reviewer. Analyze the provided task, generated code, test code, and test results. "
        "Identify issues, propose fixes, and revise the plan if necessary. "
        "If the tests pass and the code meets the requirements, indicate that the task is complete. "
        "Output your reflection as a JSON object with 'thought', 'new_plan_steps' (list of strings), and 'is_complete' (boolean)."
    )
    reflection_chain = reflection_prompt | llm.with_structured_output(Reflection)

    try:
        reflection_response = reflection_chain.invoke({
            "chat_history": state["chat_history"],
            "task": state["task"],
            "current_state": f"Current Plan:
{state['plan']}

Generated Code:
{state['code']}

Test Code:
{state['test_code']}

Test Results:
{state['test_results']}"
        })

        logger.info(f"Reflection:
Thought: {reflection_response.thought}
New Plan: {reflection_response.new_plan_steps}
Is Complete: {reflection_response.is_complete}")
        
        new_plan_str = "
".join(f"- {step}" for step in reflection_response.new_plan_steps)
        
        return {
            "plan": new_plan_str if reflection_response.new_plan_steps else state["plan"],
            "iterations": state["iterations"] + 1,
            "chat_history": state["chat_history"] + [AIMessage(content=f"Reflection: {reflection_response.thought}
New Plan: {new_plan_str}
Complete: {reflection_response.is_complete}")],
            "output": "Code generation complete, ready for PR." if reflection_response.is_complete else None
        }
    except Exception as e:
        logger.exception(f"Error in reflect_on_results: {e}")
        return {"output": f"Failed to reflect: {e}"}

def should_continue_coding(state: AgentState) -> str:
    """Conditional edge to decide whether to continue coding or proceed to GitHub."""
    if state["chat_history"] and isinstance(state["chat_history"][-1], AIMessage):
        try:
            reflection_content = state["chat_history"][-1].content
            if "Complete: True" in reflection_content:
                logger.info("Agent reflection indicates task is complete. Proceeding to GitHub.")
                return "github_integration"
        except Exception as e:
            logger.warning(f"Could not parse last chat history message for completion status: {e}")

    if state["iterations"] >= state["max_iterations"]:
        logger.warning(f"Max iterations ({state['max_iterations']}) reached. Aborting coding loop.")
        return "end_failure"
    
    logger.info(f"Continuing coding loop. Iteration {state['iterations']}/{state['max_iterations']}")
    return "continue_coding"

def create_agent_graph(openai_api_key: str):
    """
    Creates and compiles the LangGraph workflow for the autonomous coding agent.
    """
    workflow = StateGraph(AgentState)

    workflow.add_node("plan_task", plan_task)
    workflow.add_node("generate_code", generate_code)
    workflow.add_node("generate_tests", generate_tests)
    workflow.add_node("execute_tests", execute_tests)
    workflow.add_node("reflect_on_results", reflect_on_results)
    
    workflow.set_entry_point("plan_task")

    workflow.add_edge("plan_task", "generate_code")
    workflow.add_edge("generate_code", "generate_tests")
    workflow.add_edge("generate_tests", "execute_tests")
    workflow.add_edge("execute_tests", "reflect_on_results")
    
    workflow.add_conditional_edges(
        "reflect_on_results",
        should_continue_coding,
        {
            "continue_coding": "generate_code",
            "github_integration": END,
            "end_failure": END
        }
    )

    logger.info("LangGraph workflow initialized and compiled.")
    return workflow.compile(config={"configurable": {"openai_api_key": openai_api_key}})

# github_utils.py
import os
import logging
from github import Github, Auth
from github.GithubException import GithubException, UnknownObjectException
from tenacity import retry, stop_after_attempt, wait_exponential
from typing import Dict, Any, List

logger = logging.getLogger(__name__)

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def _get_github_client() -> Github:
    """Initializes and returns a GitHub client."""
    github_token = os.environ.get("GITHUB_TOKEN")
    if not github_token:
        logger.error("GITHUB_TOKEN environment variable not set.")
        raise ValueError("GITHUB_TOKEN environment variable not set.")
    
    try:
        auth = Auth.Token(github_token)
        return Github(auth=auth)
    except Exception as e:
        logger.error(f"Failed to initialize GitHub client: {e}")
        raise

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def create_branch_commit_pr(repo_owner: str, repo_name: str, task_description: str, code_content: str, test_content: str) -> str:
    """
    Creates a new branch, commits the code and tests, and opens a pull request.
    Returns the PR URL.
    """
    g = _get_github_client()
    try:
        repo = g.get_user(repo_owner).get_repo(repo_name)
    except UnknownObjectException:
        try:
            repo = g.get_organization(repo_owner).get_repo(repo_name)
        except UnknownObjectException:
            logger.error(f"Repository {repo_owner}/{repo_name} not found.")
            raise ValueError(f"Repository {repo_owner}/{repo_name} not found.")
    except GithubException as e:
        logger.error(f"Error accessing GitHub repository {repo_owner}/{repo_name}: {e}")
        raise

    branch_name = f"feature/agent-task-{hash(task_description) % 10000}" # Simple unique branch name
    commit_message = f"feat: Implement task: {task_description[:50]}..."
    pr_title = f"feat: Autonomous Agent Implements: {task_description[:70]}..."
    pr_body = (
        f"This PR was automatically generated by an autonomous coding agent to address the task:

"
        f"**Task Description:** {task_description}

"
        f"The agent generated the code and tests, executed them successfully in a sandbox, and is now proposing these changes.

"
        f"Please review the changes. If further modifications are needed, the agent can attempt to address comments."
    )

    try:
        default_branch = repo.get_branch(repo.default_branch)
        
        ref = repo.create_git_ref(ref=f"refs/heads/{branch_name}", sha=default_branch.commit.sha)
        logger.info(f"Created new branch: {branch_name}")

        latest_commit_sha = ref.object.sha

        code_blob = repo.create_git_blob(code_content, "utf-8")
        code_element = repo.create_input_git_tree_element("src/generated_code.py", "100644", "blob", code_blob.sha)

        test_blob = repo.create_git_blob(test_content, "utf-8")
        test_element = repo.create_input_git_tree_element("tests/test_generated_code.py", "100644", "blob", test_blob.sha)
        
        base_tree = repo.get_git_tree(latest_commit_sha)

        new_tree = repo.create_git_tree([code_element, test_element], base_tree)
        logger.info("Created new git tree with code and test files.")

        parent_commit = repo.get_git_commit(latest_commit_sha)
        new_commit = repo.create_git_commit(commit_message, new_tree, [parent_commit])
        logger.info(f"Created new commit: {new_commit.sha}")

        ref.edit(sha=new_commit.sha, force=False)
        logger.info(f"Branch '{branch_name}' updated to new commit.")

        pr = repo.create_pull(
            title=pr_title,
            body=pr_body,
            head=branch_name,
            base=repo.default_branch
        )
        logger.info(f"Created Pull Request: {pr.html_url}")
        return pr.html_url
    except GithubException as e:
        logger.error(f"GitHub API error during PR creation: {e.data}")
        raise
    except Exception as e:
        logger.exception(f"An unexpected error occurred during PR creation: {e}")
        raise

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def get_pr_comments(repo_owner: str, repo_name: str, pr_number: int) -> List[Dict[str, Any]]:
    """
    Fetches comments for a given pull request.
    """
    g = _get_github_client()
    try:
        repo = g.get_user(repo_owner).get_repo(repo_name)
        pr = repo.get_pull(pr_number)
        comments = [{"user": c.user.login, "body": c.body, "created_at": c.created_at.isoformat()} for c in pr.get_comments()]
        logger.info(f"Fetched {len(comments)} comments for PR #{pr_number}.")
        return comments
    except GithubException as e:
        logger.error(f"GitHub API error fetching PR comments: {e.data}")
        raise
    except Exception as e:
        logger.exception(f"An unexpected error occurred fetching PR comments: {e}")
        raise

# agent.py (updated again)
from typing import TypedDict, List, Annotated, Callable
import operator
import logging
import json
import os
from langgraph.graph import StateGraph, END
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from tenacity import retry, stop_after_attempt, wait_exponential

from tools import execute_python_code, write_file_to_sandbox, read_file_from_sandbox, list_files_in_sandbox
from github_utils import create_branch_commit_pr, get_pr_comments

logger = logging.getLogger(__name__)

class AgentState(TypedDict):
    """
    Represents the state of our agent.

    Attributes:
        task: The original task description.
        plan: The current high-level plan for execution.
        code: The generated code.
        test_code: The generated test code.
        test_results: Results from running tests.
        pull_request_url: URL of the created pull request.
        chat_history: A list of messages representing the conversation history.
        iterations: Number of self-correction iterations.
        max_iterations: Maximum allowed self-correction iterations.
        output: The final output message from the agent.
        pr_number: The number of the created pull request (for monitoring).
        github_repo_owner: Owner of the GitHub repository.
        github_repo_name: Name of the GitHub repository.
    """
    task: str
    plan: str
    code: str
    test_code: str
    test_results: str
    pull_request_url: str
    chat_history: Annotated[List[BaseMessage], operator.add]
    iterations: int
    max_iterations: int
    output: str
    pr_number: int
    github_repo_owner: str
    github_repo_name: str

class Plan(BaseModel):
    """Plan to write and test code."""
    steps: List[str] = Field(description="List of steps to complete the task.")

class CodeResponse(BaseModel):
    """Response containing generated code."""
    code: str = Field(description="The generated Python code.")

class TestResponse(BaseModel):
    """Response containing generated test code."""
    test_code: str = Field(description="The generated Python test code.")

class Reflection(BaseModel):
    """Reflection on test results and plan for next steps."""
    thought: str = Field(description="Detailed thought process on why tests failed or what needs improvement.")
    new_plan_steps: List[str] = Field(description="Revised plan steps based on reflection.")
    is_complete: bool = Field(description="True if the task is considered complete and ready for PR.")

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def _get_llm(openai_api_key: str):
    """Initializes and returns the OpenAI ChatModel."""
    return ChatOpenAI(model="gpt-4o", temperature=0.7, api_key=openai_api_key)

def _get_prompt(system_message: str):
    """Creates a chat prompt template."""
    return ChatPromptTemplate.from_messages([
        ("system", system_message),
        ("placeholder", "{chat_history}"),
        ("human", "{task}"),
        ("human", "{current_state}"), # Optional: pass current state for more context
    ])

def plan_task(state: AgentState, config: dict) -> AgentState:
    """Agent node to generate a plan for the coding task."""
    logger.info("--- Planning Task ---")
    llm = _get_llm(config["configurable"]["openai_api_key"])
    
    plan_prompt = _get_prompt(
        "You are an expert software engineer. Given a task, create a detailed plan with specific steps to implement the required code and unit tests. "
        "Break down complex problems into smaller, manageable steps. Focus on implementation details, testing strategy, and potential edge cases. "
        "Output your plan as a JSON object with a 'steps' key containing a list of strings."
    ).partial(current_state="")

    plan_chain = plan_prompt | llm.with_structured_output(Plan)
    
    try:
        plan_obj = plan_chain.invoke({
            "chat_history": state["chat_history"],
            "task": state["task"]
        })
        plan_str = "
".join(f"- {step}" for step in plan_obj.steps)
        logger.info(f"Generated Plan:
{plan_str}")

        github_repo_owner = os.environ.get("GITHUB_REPO_OWNER")
        github_repo_name = os.environ.get("GITHUB_REPO_NAME")
        if not github_repo_owner or not github_repo_name:
            logger.error("GITHUB_REPO_OWNER or GITHUB_REPO_NAME not set.")
            raise ValueError("GitHub repository details not configured.")

        return {
            "plan": plan_str,
            "chat_history": [HumanMessage(content=state["task"]), AIMessage(content=f"Plan: {plan_str}")],
            "iterations": 0,
            "max_iterations": 5,
            "github_repo_owner": github_repo_owner,
            "github_repo_name": github_repo_name
        }
    except Exception as e:
        logger.exception(f"Error in plan_task: {e}")
        return {"output": f"Failed to plan task: {e}"}

def generate_code(state: AgentState, config: dict) -> AgentState:
    """Agent node to generate code based on the plan."""
    logger.info("--- Generating Code ---")
    llm = _get_llm(config["configurable"]["openai_api_key"])
    
    code_prompt = _get_prompt(
        "You are an expert Python programmer. Given the task and current plan, write the complete Python code. "
        "Focus on writing clean, correct, and robust code. Do not include test code. "
        "Output ONLY the Python code, ensuring it's syntactically correct and ready to be executed. "
        "Wrap your code in a JSON object with a 'code' key."
    )
    code_chain = code_prompt | llm.with_structured_output(CodeResponse)

    try:
        code_response = code_chain.invoke({
            "chat_history": state["chat_history"],
            "task": state["task"],
            "current_state": f"Current Plan:
{state['plan']}"
        })
        logger.info(f"Generated Code:
{code_response.code}")
        return {
            "code": code_response.code,
            "chat_history": state["chat_history"] + [AIMessage(content=f"Generated Code:
```python
{code_response.code}
```")]
        }
    except Exception as e:
        logger.exception(f"Error in generate_code: {e}")
        return {"output": f"Failed to generate code: {e}"}

def generate_tests(state: AgentState, config: dict) -> AgentState:
    """Agent node to generate unit tests for the generated code."""
    logger.info("--- Generating Tests ---")
    llm = _get_llm(config["configurable"]["openai_api_key"])
    
    test_prompt = _get_prompt(
        "You are an expert Python developer focused on test-driven development. "
        "Given the task and the generated code, write comprehensive unit tests. "
        "Ensure tests cover functionality, edge cases, and error conditions. "
        "Use `unittest` or `pytest` if applicable, but for simple functions, direct assertions are fine. "
        "Output ONLY the Python test code. It should be executable and import the main code if necessary. "
        "Wrap your test code in a JSON object with a 'test_code' key."
    )
    test_chain = test_prompt | llm.with_structured_output(TestResponse)

    try:
        test_response = test_chain.invoke({
            "chat_history": state["chat_history"],
            "task": state["task"],
            "current_state": f"Current Plan:
{state['plan']}

Generated Code:
{state['code']}"
        })
        logger.info(f"Generated Test Code:
{test_response.test_code}")
        return {
            "test_code": test_response.test_code,
            "chat_history": state["chat_history"] + [AIMessage(content=f"Generated Tests:
```python
{test_response.test_code}
```")]
        }
    except Exception as e:
        logger.exception(f"Error in generate_tests: {e}")
        return {"output": f"Failed to generate tests: {e}"}

def execute_tests(state: AgentState) -> AgentState:
    """Agent node to execute generated code and tests in the E2B sandbox."""
    logger.info("--- Executing Code and Tests ---")
    
    write_code_result = write_file_to_sandbox("src/generated_code.py", state["code"])
    if write_code_result["status"] == "error":
        return {"output": f"Failed to write src/generated_code.py: {write_code_result['message']}"}

    write_test_result = write_file_to_sandbox("tests/test_generated_code.py", state["test_code"])
    if write_test_result["status"] == "error":
        return {"output": f"Failed to write tests/test_generated_code.py: {write_test_result['message']}"}

    # Ensure the sandbox has access to the main code for testing
    # The command assumes 'src' and 'tests' are in the root of the sandbox
    test_execution_result = execute_python_code(
        "import sys; sys.path.append('src'); import unittest; from tests.test_generated_code import *; unittest.main(exit=False)"
    )
    
    if test_execution_result["status"] == "runtime_error":
        logger.error(f"Sandbox runtime error during test execution: {test_execution_result['stderr']}")
        return {
            "test_results": f"Sandbox runtime error: {test_execution_result['stderr']}",
            "chat_history": state["chat_history"] + [AIMessage(content=f"Test Execution Failed (Sandbox Error):
{test_execution_result['stderr']}")]
        }
    
    stdout = test_execution_result["stdout"]
    stderr = test_execution_result["stderr"]
    exit_code = test_execution_result["exit_code"]

    full_output = f"Stdout:
{stdout}

Stderr:
{stderr}

Exit Code: {exit_code}"
    logger.info(f"Test Execution Result:
{full_output}")

    return {
        "test_results": full_output,
        "chat_history": state["chat_history"] + [AIMessage(content=f"Test Execution Result:
```
{full_output}
```")]
    }

def reflect_on_results(state: AgentState, config: dict) -> AgentState:
    """Agent node to reflect on test results and potentially update the plan for self-correction."""
    logger.info("--- Reflecting on Results ---")
    llm = _get_llm(config["configurable"]["openai_api_key"])
    
    reflection_prompt = _get_prompt(
        "You are an expert debugger and code reviewer. Analyze the provided task, generated code, test code, and test results. "
        "Identify issues, propose fixes, and revise the plan if necessary. "
        "If the tests pass and the code meets the requirements, indicate that the task is complete. "
        "Output your reflection as a JSON object with 'thought', 'new_plan_steps' (list of strings), and 'is_complete' (boolean)."
    )
    reflection_chain = reflection_prompt | llm.with_structured_output(Reflection)

    try:
        reflection_response = reflection_chain.invoke({
            "chat_history": state["chat_history"],
            "task": state["task"],
            "current_state": f"Current Plan:
{state['plan']}

Generated Code:
{state['code']}

Test Code:
{state['test_code']}

Test Results:
{state['test_results']}"
        })

        logger.info(f"Reflection:
Thought: {reflection_response.thought}
New Plan: {reflection_response.new_plan_steps}
Is Complete: {reflection_response.is_complete}")
        
        new_plan_str = "
".join(f"- {step}" for step in reflection_response.new_plan_steps)
        
        return {
            "plan": new_plan_str if reflection_response.new_plan_steps else state["plan"],
            "iterations": state["iterations"] + 1,
            "chat_history": state["chat_history"] + [AIMessage(content=f"Reflection: {reflection_response.thought}
New Plan: {new_plan_str}
Complete: {reflection_response.is_complete}")],
            "output": "Code generation complete, ready for PR." if reflection_response.is_complete else None
        }
    except Exception as e:
        logger.exception(f"Error in reflect_on_results: {e}")
        return {"output": f"Failed to reflect: {e}"}

def create_github_pr_node(state: AgentState) -> AgentState:
    """Agent node to create a GitHub Pull Request."""
    logger.info("--- Creating GitHub Pull Request ---")
    try:
        pr_url = create_branch_commit_pr(
            repo_owner=state["github_repo_owner"],
            repo_name=state["github_repo_name"],
            task_description=state["task"],
            code_content=state["code"],
            test_content=state["test_code"]
        )
        logger.info(f"Pull Request created: {pr_url}")
        return {
            "pull_request_url": pr_url,
            "chat_history": state["chat_history"] + [AIMessage(content=f"GitHub PR created: {pr_url}")],
            "output": f"Task completed. PR created: {pr_url}"
        }
    except Exception as e:
        logger.exception(f"Error creating GitHub PR: {e}")
        return {"output": f"Failed to create GitHub PR: {e}"}

def should_continue_coding(state: AgentState) -> str:
    """Conditional edge to decide whether to continue coding or proceed to GitHub."""
    if state["chat_history"] and isinstance(state["chat_history"][-1], AIMessage):
        try:
            reflection_content = state["chat_history"][-1].content
            if "Complete: True" in reflection_content:
                logger.info("Agent reflection indicates task is complete. Proceeding to GitHub.")
                return "create_github_pr"
        except Exception as e:
            logger.warning(f"Could not parse last chat history message for completion status: {e}")

    if state["iterations"] >= state["max_iterations"]:
        logger.warning(f"Max iterations ({state['max_iterations']}) reached. Aborting coding loop.")
        return "end_failure"
    
    logger.info(f"Continuing coding loop. Iteration {state['iterations']}/{state['max_iterations']}")
    return "continue_coding"

def create_agent_graph(openai_api_key: str):
    """
    Creates and compiles the LangGraph workflow for the autonomous coding agent.
    """
    workflow = StateGraph(AgentState)

    workflow.add_node("plan_task", plan_task)
    workflow.add_node("generate_code", generate_code)
    workflow.add_node("generate_tests", generate_tests)
    workflow.add_node("execute_tests", execute_tests)
    workflow.add_node("reflect_on_results", reflect_on_results)
    workflow.add_node("create_github_pr", create_github_pr_node)
    
    workflow.set_entry_point("plan_task")

    workflow.add_edge("plan_task", "generate_code")
    workflow.add_edge("generate_code", "generate_tests")
    workflow.add_edge("generate_tests", "execute_tests")
    workflow.add_edge("execute_tests", "reflect_on_results")
    
    workflow.add_conditional_edges(
        "reflect_on_results",
        should_continue_coding,
        {
            "continue_coding": "generate_code",
            "create_github_pr": "create_github_pr",
            "end_failure": END
        }
    )
    workflow.add_edge("create_github_pr", END)

    logger.info("LangGraph workflow initialized and compiled.")
    return workflow.compile(config={"configurable": {"openai_api_key": openai_api_key}})

# main.py (updated for LangSmith note)
import os
import logging
from dotenv import load_dotenv
from typing import TypedDict, List, Annotated
from langgraph.graph import StateGraph, END
from agent import create_agent_graph, AgentState

load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Optional: Configure LangSmith for observability
# Set LANGCHAIN_TRACING_V2="true" and LANGCHAIN_API_KEY="your_langsmith_api_key" in your .env
# LANGCHAIN_PROJECT should also be set, e.g., "Autonomous Coding Agent"
# For local development, you might not need to set these explicitly if you're not using LangSmith.
# However, for production, these are crucial.

def run_agent(task: str) -> str:
    """
    Initializes and runs the autonomous coding agent with a given task.
    """
    openai_api_key = os.environ.get("OPENAI_API_KEY")
    e2b_api_key = os.environ.get("E2B_API_KEY")
    github_token = os.environ.get("GITHUB_TOKEN")
    github_repo_owner = os.environ.get("GITHUB_REPO_OWNER")
    github_repo_name = os.environ.get("GITHUB_REPO_NAME")

    if not all([openai_api_key, e2b_api_key, github_token, github_repo_owner, github_repo_name]):
        logger.error("Missing one or more required environment variables.")
        raise ValueError("Ensure OPENAI_API_KEY, E2B_API_KEY, GITHUB_TOKEN, GITHUB_REPO_OWNER, GITHUB_REPO_NAME are set.")

    workflow = create_agent_graph(openai_api_key)
    
    inputs: AgentState = {
        "task": task, 
        "chat_history": [],
        "plan": "",
        "code": "",
        "test_code": "",
        "test_results": "",
        "pull_request_url": "",
        "iterations": 0,
        "max_iterations": 5, # Default, can be overridden by plan_task
        "output": "",
        "pr_number": 0, # Placeholder
        "github_repo_owner": github_repo_owner,
        "github_repo_name": github_repo_name
    }
    
    logger.info(f"Starting agent for task: {task}")
    final_output = None
    try:
        # The LangGraph stream method automatically integrates with LangSmith if environment variables are set.
        for s in workflow.stream(inputs, {"configurable": {"thread_id": "1"}}):
            logger.info(f"Agent step: {s}")
            final_output = s
        
        if final_output:
            # final_output is {"node_name": state_dict} from the last streamed step
            last_state = next(iter(final_output.values()), {}) if isinstance(final_output, dict) else {}
            return last_state.get("output", str(final_output))
        else:
            return "Agent finished without explicit output."

    except Exception as e:
        logger.exception(f"An error occurred during agent execution: {e}")
        return f"Agent execution failed: {e}"

if __name__ == "__main__":
    test_task = "Create a Python function `multiply(a, b)` that returns the product of two numbers. Include a docstring and a simple unit test for it. The function should be in a file named `src/generated_code.py` and tests in `tests/test_generated_code.py`."
    print("--- Running autonomous coding agent ---")
    result = run_agent(test_task)
    print("
--- Agent Finished ---")
    print(f"Final Result: {result}")