aurelian 0.3.2__py3-none-any.whl

aurelian/agents/filesystem/filesystem_mcp.py

@@ -0,0 +1,89 @@
+"""
+MCP tools for filesystem operations.
+"""
+import os
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+
+import aurelian.agents.filesystem.filesystem_tools as ft
+from aurelian.dependencies.workdir import WorkDir, HasWorkdir
+from pydantic_ai import RunContext
+
+# Initialize FastMCP server
+mcp = FastMCP("filesystem", instructions="Filesystem operations")
+
+
+def deps() -> HasWorkdir:
+    deps = HasWorkdir()
+    # Set the location from environment variable or default
+    loc = os.getenv("AURELIAN_WORKDIR", "/tmp/aurelian")
+    deps.workdir = WorkDir(loc)
+    return deps
+
+def ctx() -> RunContext[HasWorkdir]:
+    rc: RunContext[HasWorkdir] = RunContext[HasWorkdir](
+        deps=deps(),
+        model=None, usage=None, prompt=None,
+    )
+    return rc
+
+
+@mcp.tool()
+async def inspect_file(data_file: str) -> str:
+    """
+    Inspect a file in the working directory.
+
+    Args:
+        data_file: name of file
+
+    Returns:
+        Contents of the file
+    """
+    return await ft.inspect_file(ctx(), data_file)
+
+
+@mcp.tool()
+async def download_url_as_markdown(url: str, local_file_name: str) -> ft.DownloadResult:
+    """
+    Download a URL and convert to markdown.
+
+    Args:
+        url: The URL to download
+        local_file_name: The name to save the file as
+
+    Returns:
+        DownloadResult with file name and number of lines
+    """
+    return await ft.download_url_as_markdown(ctx(), url, local_file_name)
+
+
+@mcp.tool()
+async def list_files() -> str:
+    """
+    List files in the working directory.
+
+    Returns:
+        A string listing the files in the working directory
+    """
+    return await ft.list_files(ctx())
+
+
+@mcp.tool()
+async def write_to_file(file_name: str, data: str) -> str:
+    """
+    Write data to a file in the working directory.
+
+    Args:
+        file_name: Name of file to write
+        data: Data to write to the file
+
+    Returns:
+        Confirmation message
+    """
+    return await ft.write_to_file(ctx(), file_name, data)
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport='stdio')

aurelian/agents/filesystem/filesystem_tools.py

@@ -0,0 +1,95 @@
+import re
+from typing import Any
+
+import httpx
+from pydantic import BaseModel
+from pydantic_ai import RunContext
+
+from aurelian.dependencies.workdir import WorkDir, HasWorkdir
+
+
+class DownloadResult(BaseModel):
+    file_name: str
+    num_lines: int
+
+async def inspect_file(ctx: RunContext[HasWorkdir], data_file: str) -> str:
+    """
+    Inspect a file in the working directory.
+
+    Args:
+        ctx:
+        data_file: name of file
+
+    Returns:
+
+    """
+    print(f"Inspecting file: {data_file}")
+    return ctx.deps.workdir.read_file(data_file)
+
+
+async def download_url_as_markdown(ctx: RunContext[HasWorkdir], url: str, local_file_name: str) -> DownloadResult:
+    """
+    Download contents of a web page.
+
+    Args:
+        ctx: context
+        url: URL of the web page
+        local_file_name: Name of the local file to save the
+
+    Returns:
+        DownloadResult: information about the downloaded file
+    """
+    workdir: WorkDir = ctx.deps.workdir
+    from markdownify import markdownify
+
+    async with httpx.AsyncClient() as client:
+        response = await client.get(url, timeout=20.0)
+        response.raise_for_status()  # Raise an exception for bad status codes
+
+        # Convert the HTML content to Markdown
+        markdown_content = markdownify(response.text).strip()
+
+        # Remove multiple line breaks
+        markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
+
+        # Assuming write_file is also async
+        workdir.write_file(local_file_name, markdown_content)
+
+        return DownloadResult(file_name=local_file_name, num_lines=len(markdown_content.split("\n")))
+
+
+async def list_files(ctx: RunContext[HasWorkdir]) -> str:
+    """
+    List files in the working directory.
+
+    Args:
+        ctx:
+
+    Returns:
+
+    """
+    return "\n".join(ctx.deps.workdir.list_file_names())
+
+
+async def write_to_file(ctx: RunContext[HasWorkdir], file_name: str, data: str) -> str:
+    """
+    Write data to a file in the working directory.
+
+    Args:
+        ctx:
+        file_name:
+        data:
+
+    Returns:
+
+    """
+    print(f"Writing data to file: {file_name}")
+    ctx.deps.workdir.write_file(file_name, data)
+    return f"Data written to {file_name}"
+
+
+async def show_local_files(ctx: RunContext[HasWorkdir]) -> str:
+    file_names = ctx.deps.workdir.list_file_names()
+    if file_names:
+        return f"Local files: {file_names}"
+    return "No files currently in the working directory"

aurelian/agents/github/github_agent.py

@@ -0,0 +1,83 @@
+"""
+Agent for working with GitHub repositories, issues, and pull requests.
+"""
+from pydantic import BaseModel
+from typing import List, Optional, Dict, Any, Union
+
+from pydantic_ai import Agent, Tool, RunContext
+
+from aurelian.agents.github.github_config import GitHubDependencies
+from aurelian.agents.github.github_tools import (
+    list_pull_requests,
+    view_pull_request,
+    list_issues,
+    view_issue,
+    get_pr_closing_issues,
+    get_commit_before_pr,
+    search_code,
+    clone_repository
+)
+from aurelian.agents.filesystem.filesystem_tools import inspect_file, list_files
+
+
+class GitHubRepoSummary(BaseModel):
+    """
+    Summary of a GitHub repository.
+    """
+    repo_name: str
+    owner: str
+    description: Optional[str] = None
+    stars: Optional[int] = None
+    forks: Optional[int] = None
+    open_issues: Optional[int] = None
+    open_prs: Optional[int] = None
+    
+    
+SYSTEM = """
+You are a GitHub repository assistant with access to GitHub data through the GitHub CLI.
+
+I can help you with:
+- Listing and viewing pull requests
+- Examining issues and their status
+- Finding what issues are closed by pull requests
+- Searching code within repositories
+- Cloning repositories locally
+- Examining commit history, especially commits before a PR was created
+
+When answering questions about GitHub repositories, I'll use the GitHub CLI to get 
+the most up-to-date information. If you want to work with a specific repository, 
+I can clone it and then work with the files directly.
+
+You can ask me to:
+- List recent PRs or issues
+- View details of specific PRs or issues
+- Find what issues a PR closes
+- Search for code patterns in repositories
+- Clone repositories for deeper analysis
+- Identify the commit before a PR branch was created
+
+I'm best at answering specific questions that make use of GitHub's structure, 
+for example finding connections between PRs and issues, or examining code across 
+a repository.
+"""
+
+
+github_tools = [
+    Tool(list_pull_requests),
+    Tool(view_pull_request),
+    Tool(list_issues),
+    Tool(view_issue),
+    Tool(get_pr_closing_issues),
+    Tool(get_commit_before_pr),
+    Tool(search_code),
+    Tool(clone_repository),
+    Tool(inspect_file),
+    Tool(list_files),
+]
+
+github_agent = Agent(
+    model="openai:gpt-4o",
+    deps_type=GitHubDependencies,
+    system_prompt=SYSTEM,
+    tools=github_tools,
+)

aurelian/agents/github/github_cli.py

@@ -0,0 +1,248 @@
+"""
+Command-line interface for the GitHub agent.
+"""
+import os
+import sys
+import json
+import click
+import tempfile
+import logging
+from pathlib import Path
+from typing import Optional
+
+from aurelian.agents.github.github_agent import github_agent
+from aurelian.agents.github.github_config import GitHubDependencies
+from aurelian.dependencies.workdir import WorkDir
+
+
+def setup_logging():
+    """Set up logging for the CLI."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        handlers=[logging.StreamHandler(sys.stdout)]
+    )
+
+
+def create_deps(workdir: Optional[str] = None) -> GitHubDependencies:
+    """
+    Create dependencies for the GitHub agent.
+    
+    Args:
+        workdir: Working directory path. If None, uses current directory.
+        
+    Returns:
+        GitHubDependencies instance
+    """
+    deps = GitHubDependencies()
+    if workdir:
+        wd_path = Path(workdir)
+        # Create directory if it doesn't exist
+        if not wd_path.exists():
+            wd_path.mkdir(parents=True)
+    else:
+        workdir = os.getcwd()
+    
+    deps.workdir = WorkDir(location=workdir)
+    return deps
+
+
+@click.group()
+@click.option('--workdir', '-w', help='Working directory path')
+@click.pass_context
+def cli(ctx, workdir):
+    """GitHub repository assistant CLI."""
+    setup_logging()
+    ctx.obj = create_deps(workdir)
+
+
+@cli.command()
+@click.argument('repo')
+@click.option('--directory', '-d', help='Directory name to clone into')
+@click.option('--branch', '-b', help='Branch to clone')
+@click.option('--depth', type=int, help='Depth limit for clone (for shallow clones)')
+@click.pass_obj
+def clone(deps, repo, directory, branch, depth):
+    """Clone a GitHub repository."""
+    from aurelian.agents.github.github_tools import clone_repository
+    from pydantic_ai import RunContext
+    
+    ctx = RunContext[GitHubDependencies](deps=deps, model="openai:gpt-4o", usage=None, prompt="")
+    
+    try:
+        clone_dir = clone_repository(ctx, repo, directory, branch, depth)
+        click.echo(f"Successfully cloned {repo} to {clone_dir}")
+        
+        # Optionally update workdir to the cloned repo
+        if click.confirm("Do you want to set the working directory to the cloned repository?"):
+            deps.workdir = WorkDir(location=clone_dir)
+            click.echo(f"Working directory updated to: {clone_dir}")
+    except Exception as e:
+        click.echo(f"Error cloning repository: {str(e)}", err=True)
+        sys.exit(1)
+
+
+@cli.command()
+@click.option('--state', '-s', default='open', help='Filter by state (open, closed, merged, all)')
+@click.option('--limit', '-l', default=10, type=int, help='Maximum number of PRs to return')
+@click.option('--label', help='Filter by label')
+@click.option('--author', '-a', help='Filter by author')
+@click.option('--base', help='Filter by base branch')
+@click.option('--repo', '-r', help='Repository (format: owner/repo)')
+@click.option('--output', '-o', help='Output file path')
+@click.option('--format', 'format_', type=click.Choice(['pretty', 'json']), default='pretty', 
+              help='Output format (pretty or json)')
+@click.pass_obj
+def list_prs(deps, state, limit, label, author, base, repo, output, format_):
+    """List pull requests from GitHub."""
+    from aurelian.agents.github.github_tools import list_pull_requests
+    from pydantic_ai import RunContext
+    
+    ctx = RunContext[GitHubDependencies](deps=deps, model="openai:gpt-4o", usage=None, prompt="")
+    
+    try:
+        results = list_pull_requests(ctx, state, limit, label, author, base, repo)
+        
+        if format_ == 'json':
+            output_data = json.dumps(results, indent=2)
+        else:
+            output_data = "Pull Requests:\n"
+            for i, pr in enumerate(results, 1):
+                output_data += f"{i}. #{pr['number']} - {pr['title']} ({pr['state']})\n"
+                output_data += f"   Author: {pr['author']['login']}\n"
+                output_data += f"   Created: {pr['createdAt']}\n"
+                output_data += f"   URL: {pr['url']}\n\n"
+        
+        if output:
+            with open(output, 'w') as f:
+                f.write(output_data)
+            click.echo(f"Results written to {output}")
+        else:
+            click.echo(output_data)
+    except Exception as e:
+        click.echo(f"Error listing pull requests: {str(e)}", err=True)
+        sys.exit(1)
+
+
+@cli.command()
+@click.argument('pr_number')
+@click.option('--repo', '-r', help='Repository (format: owner/repo)')
+@click.option('--output', '-o', help='Output file path')
+@click.option('--format', 'format_', type=click.Choice(['pretty', 'json']), default='pretty',
+              help='Output format (pretty or json)')
+@click.pass_obj
+def view_pr(deps, pr_number, repo, output, format_):
+    """View a specific pull request."""
+    from aurelian.agents.github.github_tools import view_pull_request
+    from pydantic_ai import RunContext
+    
+    ctx = RunContext[GitHubDependencies](deps=deps, model="openai:gpt-4o", usage=None, prompt="")
+    
+    try:
+        result = view_pull_request(ctx, pr_number, repo)
+        
+        if format_ == 'json':
+            output_data = json.dumps(result, indent=2)
+        else:
+            output_data = f"PR #{result['number']}: {result['title']}\n"
+            output_data += f"State: {result['state']}\n"
+            output_data += f"Author: {result['author']['login']}\n"
+            output_data += f"Created: {result['createdAt']}\n"
+            output_data += f"URL: {result['url']}\n"
+            output_data += f"Base: {result['baseRefName']} <- Head: {result['headRefName']}\n\n"
+            output_data += f"Description:\n{result['body']}\n\n"
+            
+            if result.get('closingIssues'):
+                output_data += "Closes Issues:\n"
+                for issue in result['closingIssues']:
+                    output_data += f"- #{issue['number']}: {issue['title']}\n"
+        
+        if output:
+            with open(output, 'w') as f:
+                f.write(output_data)
+            click.echo(f"Results written to {output}")
+        else:
+            click.echo(output_data)
+    except Exception as e:
+        click.echo(f"Error viewing pull request: {str(e)}", err=True)
+        sys.exit(1)
+
+
+@cli.command()
+@click.argument('pr_number')
+@click.option('--repo', '-r', help='Repository (format: owner/repo)')
+@click.pass_obj
+def commit_before_pr(deps, pr_number, repo):
+    """Get the commit before a PR branch was created."""
+    from aurelian.agents.github.github_tools import get_commit_before_pr
+    from pydantic_ai import RunContext
+    
+    ctx = RunContext[GitHubDependencies](deps=deps, model="openai:gpt-4o", usage=None, prompt="")
+    
+    try:
+        result = get_commit_before_pr(ctx, pr_number, repo)
+        click.echo(f"Commit before PR #{pr_number} was created:")
+        click.echo(f"Hash: {result['hash']}")
+        click.echo(f"Message:\n{result['message']}")
+    except Exception as e:
+        click.echo(f"Error getting commit before PR: {str(e)}", err=True)
+        sys.exit(1)
+
+
+@cli.command()
+@click.argument('query')
+@click.option('--repo', '-r', help='Repository (format: owner/repo)')
+@click.option('--output', '-o', help='Output file path')
+@click.pass_obj
+def search(deps, query, repo, output):
+    """Search for code in a repository."""
+    from aurelian.agents.github.github_tools import search_code
+    from pydantic_ai import RunContext
+    
+    ctx = RunContext[GitHubDependencies](deps=deps, model="openai:gpt-4o", usage=None, prompt="")
+    
+    try:
+        results = search_code(ctx, query, repo)
+        
+        output_data = f"Search results for: {query}\n\n"
+        for i, match in enumerate(results, 1):
+            output_data += f"{i}. {match['path']}\n"
+            for text_match in match.get('textMatches', []):
+                if 'fragment' in text_match:
+                    output_data += f"   ...{text_match['fragment']}...\n"
+            output_data += "\n"
+        
+        if output:
+            with open(output, 'w') as f:
+                f.write(output_data)
+            click.echo(f"Results written to {output}")
+        else:
+            click.echo(output_data)
+    except Exception as e:
+        click.echo(f"Error searching code: {str(e)}", err=True)
+        sys.exit(1)
+
+
+@cli.command()
+@click.option('--query', '-q', help='Initial question or prompt for the agent')
+@click.pass_obj
+def chat(deps, query):
+    """Start an interactive chat with the GitHub agent."""
+    click.echo("Starting chat with GitHub agent. Type 'exit' or 'quit' to end.")
+    
+    if query:
+        response = github_agent.chat(query, deps=deps)
+        click.echo(f"Agent: {response.content}")
+    
+    while True:
+        user_input = click.prompt("You", type=str)
+        if user_input.lower() in ('exit', 'quit'):
+            click.echo("Exiting chat.")
+            break
+        
+        response = github_agent.chat(user_input, deps=deps)
+        click.echo(f"Agent: {response.content}")
+
+
+if __name__ == '__main__':
+    cli()

aurelian/agents/github/github_config.py

@@ -0,0 +1,22 @@
+"""
+Configuration for the GitHub agent.
+"""
+import os
+
+from aurelian.dependencies.workdir import HasWorkdir, WorkDir
+
+
+class GitHubDependencies(HasWorkdir):
+    """
+    GitHub agent dependencies that include a working directory.
+    
+    This allows the agent to maintain state and access files on the local filesystem.
+    """
+    pass
+
+
+def get_config() -> GitHubDependencies:
+    workdir_path = os.environ.get("AURELIAN_WORKDIR", None)
+    workdir = WorkDir(location=workdir_path) if workdir_path else None
+
+    return GitHubDependencies(workdir=workdir)

aurelian/agents/github/github_gradio.py

@@ -0,0 +1,152 @@
+"""
+Gradio UI for the GitHub agent.
+"""
+from typing import List, Optional
+import tempfile
+import os
+import shutil
+
+import gradio as gr
+
+from aurelian.agents.github.github_config import GitHubDependencies
+from aurelian.agents.github.github_agent import github_agent
+from aurelian.utils.async_utils import run_sync
+from aurelian.dependencies.workdir import WorkDir
+
+
+def create_demo(workspace_dir: Optional[str] = None) -> gr.Blocks:
+    """
+    Create a Gradio demo for the GitHub agent.
+    
+    Args:
+        workspace_dir: Directory to use as workspace. If None, a temporary directory is created.
+        
+    Returns:
+        Gradio Blocks interface
+    """
+    temp_dir = None
+    if workspace_dir is None:
+        temp_dir = tempfile.mkdtemp()
+        workspace_dir = temp_dir
+    
+    deps = GitHubDependencies()
+    deps.workdir = WorkDir(location=workspace_dir)
+    
+    with gr.Blocks(title="GitHub Agent") as demo:
+        gr.Markdown("# GitHub Repository Agent")
+        gr.Markdown("""
+        This agent helps you work with GitHub repositories, issues, and pull requests.
+        
+        You can:
+        - List and view pull requests and issues
+        - Find connections between PRs and issues
+        - Search code in repositories
+        - Clone repositories for deeper analysis
+        - Examine commit history
+        
+        The agent uses the GitHub CLI, so you need to have it installed and authenticated.
+        """)
+        
+        current_dir_md = gr.Markdown(f"**Working Directory**: {workspace_dir}")
+        
+        def update_current_dir():
+            return f"**Working Directory**: {deps.workdir.location}"
+        
+        with gr.Row():
+            with gr.Column():
+                repo_input = gr.Textbox(label="Repository (owner/repo)")
+                clone_btn = gr.Button("Clone Repository")
+            
+            with gr.Column():
+                directory_input = gr.Textbox(label="Directory Name (optional)")
+                branch_input = gr.Textbox(label="Branch (optional)")
+        
+        def clone_repo(repo: str, directory: Optional[str] = None, branch: Optional[str] = None):
+            from aurelian.agents.github.github_tools import clone_repository
+            from pydantic_ai import RunContext
+            
+            ctx = RunContext[GitHubDependencies](
+                deps=deps, 
+                model="openai:gpt-4o", 
+                usage=None, 
+                prompt=""
+            )
+            
+            if not repo:
+                return "Please enter a repository name in the format owner/repo", update_current_dir()
+            
+            try:
+                clone_dir = clone_repository(ctx, repo, directory, branch)
+                # Set workdir to the cloned repo
+                deps.workdir = WorkDir(location=clone_dir)
+                return f"Successfully cloned {repo} to {clone_dir}", update_current_dir()
+            except Exception as e:
+                return f"Error cloning repository: {str(e)}", update_current_dir()
+        
+        clone_result = gr.Textbox(label="Clone Result")
+        clone_btn.click(clone_repo, inputs=[repo_input, directory_input, branch_input], outputs=[clone_result, current_dir_md])
+        
+        def get_info(query: str, history: List[List[str]]) -> str:
+            flattened_history = []
+            for h in history:
+                flattened_history.extend(h)
+            
+            return run_sync(lambda: github_agent.chat(
+                query, 
+                deps=deps
+            ).content)
+
+        chatbot = gr.Chatbot()
+        msg = gr.Textbox(label="Your question about GitHub repositories")
+        
+        msg.submit(get_info, inputs=[msg, chatbot], outputs=chatbot)
+        
+        examples = [
+            "List the most recent open pull requests",
+            "What issues will be closed by PR #31?",
+            "What was the commit before PR #31 was created?",
+            "Search for code related to 'RunContext'",
+            "Clone the repository 'obophenotype/cell-ontology' and list its files",
+        ]
+        gr.Examples(examples=examples, inputs=msg)
+    
+    # Cleanup temporary directory when demo is closed
+    if temp_dir:
+        def cleanup():
+            if os.path.exists(temp_dir):
+                shutil.rmtree(temp_dir)
+        
+        demo.close_callbacks.append(cleanup)
+    
+    return demo
+
+
+def chat(deps: Optional[GitHubDependencies] = None, **kwargs):
+    """
+    Initialize a chat interface for the GitHub agent.
+    
+    Args:
+        deps: Optional dependencies configuration
+        **kwargs: Additional arguments to pass to the agent
+        
+    Returns:
+        A Gradio chat interface
+    """
+    if deps is None:
+        deps = GitHubDependencies()
+        deps.workdir = WorkDir(location=os.getcwd())
+
+    def get_info(query: str, history: List[str]) -> str:
+        result = run_sync(lambda: github_agent.chat(query, deps=deps, **kwargs))
+        return result.content
+
+    return gr.ChatInterface(
+        fn=get_info,
+        title="GitHub AI Assistant",
+        examples=[
+            "List all open issues in monarch-initiative/aurelian",
+            "Show me pull request #31 in monarch-initiative/aurelian",
+            "Search for code related to RunContext in monarch-initiative/aurelian",
+            "Get the commit before PR #31 was created"
+        ]
+    )