The New Archon (Beta) - The Operating System for AI Coding Assistants!

original_archon/iterations/v1-single-agent/.env.example

@@ -0,0 +1,19 @@
+# Get your Open AI API Key by following these instructions -
+# https://help.openai.com/en/articles/4936850-where-do-i-find-my-openai-api-key
+# You only need this environment variable set if you are using GPT (and not Ollama)
+OPENAI_API_KEY=
+
+# For the Supabase version (sample_supabase_agent.py), set your Supabase URL and Service Key.
+# Get your SUPABASE_URL from the API section of your Supabase project settings -
+# https://supabase.com/dashboard/project/<your project ID>/settings/api
+SUPABASE_URL=
+
+# Get your SUPABASE_SERVICE_KEY from the API section of your Supabase project settings -
+# https://supabase.com/dashboard/project/<your project ID>/settings/api
+# On this page it is called the service_role secret.
+SUPABASE_SERVICE_KEY=
+
+# The LLM you want to use from OpenAI. See the list of models here:
+# https://platform.openai.com/docs/models
+# Example: gpt-4o-mini
+LLM_MODEL=

original_archon/iterations/v1-single-agent/README.md

@@ -0,0 +1,122 @@
+# Archon V1 - Basic Pydantic AI Agent to Build other Pydantic AI Agents
+
+This is the first iteration of the Archon project - no use of LangGraph and built with a single AI agent to keep things very simple and introductory.
+
+An intelligent documentation crawler and RAG (Retrieval-Augmented Generation) agent built using Pydantic AI and Supabase that is capable of building other Pydantic AI agents. The agent crawls the Pydantic AI documentation, stores content in a vector database, and provides Pydantic AI agent code by retrieving and analyzing relevant documentation chunks.
+
+## Features
+
+- Pydantic AI documentation crawling and chunking
+- Vector database storage with Supabase
+- Semantic search using OpenAI embeddings
+- RAG-based question answering
+- Support for code block preservation
+- Streamlit UI for interactive querying
+
+## Prerequisites
+
+- Python 3.11+
+- Supabase account and database
+- OpenAI API key
+- Streamlit (for web interface)
+
+## Installation
+
+1. Clone the repository:
+```bash
+git clone https://github.com/coleam00/archon.git
+cd archon/iterations/v1-single-agent
+```
+
+2. Install dependencies (recommended to use a Python virtual environment):
+```bash
+python -m venv venv
+source venv/bin/activate  # On Windows: venv\Scripts\activate
+pip install -r requirements.txt
+```
+
+3. Set up environment variables:
+   - Rename `.env.example` to `.env`
+   - Edit `.env` with your API keys and preferences:
+   ```env
+   OPENAI_API_KEY=your_openai_api_key
+   SUPABASE_URL=your_supabase_url
+   SUPABASE_SERVICE_KEY=your_supabase_service_key
+   LLM_MODEL=gpt-4o-mini  # or your preferred OpenAI model
+   ```
+
+## Usage
+
+### Database Setup
+
+Execute the SQL commands in `site_pages.sql` to:
+1. Create the necessary tables
+2. Enable vector similarity search
+3. Set up Row Level Security policies
+
+In Supabase, do this by going to the "SQL Editor" tab and pasting in the SQL into the editor there. Then click "Run".
+
+### Crawl Documentation
+
+To crawl and store documentation in the vector database:
+
+```bash
+python crawl_pydantic_ai_docs.py
+```
+
+This will:
+1. Fetch URLs from the documentation sitemap
+2. Crawl each page and split into chunks
+3. Generate embeddings and store in Supabase
+
+### Streamlit Web Interface
+
+For an interactive web interface to query the documentation:
+
+```bash
+streamlit run streamlit_ui.py
+```
+
+The interface will be available at `http://localhost:8501`
+
+## Configuration
+
+### Database Schema
+
+The Supabase database uses the following schema:
+```sql
+CREATE TABLE site_pages (
+    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
+    url TEXT,
+    chunk_number INTEGER,
+    title TEXT,
+    summary TEXT,
+    content TEXT,
+    metadata JSONB,
+    embedding VECTOR(1536)
+);
+```
+
+### Chunking Configuration
+
+You can configure chunking parameters in `crawl_pydantic_ai_docs.py`:
+```python
+chunk_size = 5000  # Characters per chunk
+```
+
+The chunker intelligently preserves:
+- Code blocks
+- Paragraph boundaries
+- Sentence boundaries
+
+## Project Structure
+
+- `crawl_pydantic_ai_docs.py`: Documentation crawler and processor
+- `pydantic_ai_expert.py`: RAG agent implementation
+- `streamlit_ui.py`: Web interface
+- `site_pages.sql`: Database setup commands
+- `requirements.txt`: Project dependencies
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.

original_archon/iterations/v1-single-agent/crawl_pydantic_ai_docs.py

@@ -0,0 +1,245 @@
+import os
+import sys
+import json
+import asyncio
+import requests
+from xml.etree import ElementTree
+from typing import List, Dict, Any
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from urllib.parse import urlparse
+from dotenv import load_dotenv
+
+from crawl4ai import AsyncWebCrawler, BrowserConfig, CrawlerRunConfig, CacheMode
+from openai import AsyncOpenAI
+from supabase import create_client, Client
+
+load_dotenv()
+
+# Initialize OpenAI and Supabase clients
+openai_client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+supabase: Client = create_client(
+    os.getenv("SUPABASE_URL"),
+    os.getenv("SUPABASE_SERVICE_KEY")
+)
+
+@dataclass
+class ProcessedChunk:
+    url: str
+    chunk_number: int
+    title: str
+    summary: str
+    content: str
+    metadata: Dict[str, Any]
+    embedding: List[float]
+
+def chunk_text(text: str, chunk_size: int = 5000) -> List[str]:
+    """Split text into chunks, respecting code blocks and paragraphs."""
+    chunks = []
+    start = 0
+    text_length = len(text)
+
+    while start < text_length:
+        # Calculate end position
+        end = start + chunk_size
+
+        # If we're at the end of the text, just take what's left
+        if end >= text_length:
+            chunks.append(text[start:].strip())
+            break
+
+        # Try to find a code block boundary first (```)
+        chunk = text[start:end]
+        code_block = chunk.rfind('```')
+        if code_block != -1 and code_block > chunk_size * 0.3:
+            end = start + code_block
+
+        # If no code block, try to break at a paragraph
+        elif '\n\n' in chunk:
+            # Find the last paragraph break
+            last_break = chunk.rfind('\n\n')
+            if last_break > chunk_size * 0.3:  # Only break if we're past 30% of chunk_size
+                end = start + last_break
+
+        # If no paragraph break, try to break at a sentence
+        elif '. ' in chunk:
+            # Find the last sentence break
+            last_period = chunk.rfind('. ')
+            if last_period > chunk_size * 0.3:  # Only break if we're past 30% of chunk_size
+                end = start + last_period + 1
+
+        # Extract chunk and clean it up
+        chunk = text[start:end].strip()
+        if chunk:
+            chunks.append(chunk)
+
+        # Move start position for next chunk
+        start = max(start + 1, end)
+
+    return chunks
+
+async def get_title_and_summary(chunk: str, url: str) -> Dict[str, str]:
+    """Extract title and summary using GPT-4."""
+    system_prompt = """You are an AI that extracts titles and summaries from documentation chunks.
+    Return a JSON object with 'title' and 'summary' keys.
+    For the title: If this seems like the start of a document, extract its title. If it's a middle chunk, derive a descriptive title.
+    For the summary: Create a concise summary of the main points in this chunk.
+    Keep both title and summary concise but informative."""
+    
+    try:
+        response = await openai_client.chat.completions.create(
+            model=os.getenv("LLM_MODEL", "gpt-4o-mini"),
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": f"URL: {url}\n\nContent:\n{chunk[:1000]}..."}  # Send first 1000 chars for context
+            ],
+            response_format={ "type": "json_object" }
+        )
+        return json.loads(response.choices[0].message.content)
+    except Exception as e:
+        print(f"Error getting title and summary: {e}")
+        return {"title": "Error processing title", "summary": "Error processing summary"}
+
+async def get_embedding(text: str) -> List[float]:
+    """Get embedding vector from OpenAI."""
+    try:
+        response = await openai_client.embeddings.create(
+            model="text-embedding-3-small",
+            input=text
+        )
+        return response.data[0].embedding
+    except Exception as e:
+        print(f"Error getting embedding: {e}")
+        return [0] * 1536  # Return zero vector on error
+
+async def process_chunk(chunk: str, chunk_number: int, url: str) -> ProcessedChunk:
+    """Process a single chunk of text."""
+    # Get title and summary
+    extracted = await get_title_and_summary(chunk, url)
+    
+    # Get embedding
+    embedding = await get_embedding(chunk)
+    
+    # Create metadata
+    metadata = {
+        "source": "pydantic_ai_docs",
+        "chunk_size": len(chunk),
+        "crawled_at": datetime.now(timezone.utc).isoformat(),
+        "url_path": urlparse(url).path
+    }
+    
+    return ProcessedChunk(
+        url=url,
+        chunk_number=chunk_number,
+        title=extracted['title'],
+        summary=extracted['summary'],
+        content=chunk,  # Store the original chunk content
+        metadata=metadata,
+        embedding=embedding
+    )
+
+async def insert_chunk(chunk: ProcessedChunk):
+    """Insert a processed chunk into Supabase."""
+    try:
+        data = {
+            "url": chunk.url,
+            "chunk_number": chunk.chunk_number,
+            "title": chunk.title,
+            "summary": chunk.summary,
+            "content": chunk.content,
+            "metadata": chunk.metadata,
+            "embedding": chunk.embedding
+        }
+        
+        result = supabase.table("site_pages").insert(data).execute()
+        print(f"Inserted chunk {chunk.chunk_number} for {chunk.url}")
+        return result
+    except Exception as e:
+        print(f"Error inserting chunk: {e}")
+        return None
+
+async def process_and_store_document(url: str, markdown: str):
+    """Process a document and store its chunks in parallel."""
+    # Split into chunks
+    chunks = chunk_text(markdown)
+    
+    # Process chunks in parallel
+    tasks = [
+        process_chunk(chunk, i, url) 
+        for i, chunk in enumerate(chunks)
+    ]
+    processed_chunks = await asyncio.gather(*tasks)
+    
+    # Store chunks in parallel
+    insert_tasks = [
+        insert_chunk(chunk) 
+        for chunk in processed_chunks
+    ]
+    await asyncio.gather(*insert_tasks)
+
+async def crawl_parallel(urls: List[str], max_concurrent: int = 5):
+    """Crawl multiple URLs in parallel with a concurrency limit."""
+    browser_config = BrowserConfig(
+        headless=True,
+        verbose=False,
+        extra_args=["--disable-gpu", "--disable-dev-shm-usage", "--no-sandbox"],
+    )
+    crawl_config = CrawlerRunConfig(cache_mode=CacheMode.BYPASS)
+
+    # Create the crawler instance
+    crawler = AsyncWebCrawler(config=browser_config)
+    await crawler.start()
+
+    try:
+        # Create a semaphore to limit concurrency
+        semaphore = asyncio.Semaphore(max_concurrent)
+        
+        async def process_url(url: str):
+            async with semaphore:
+                result = await crawler.arun(
+                    url=url,
+                    config=crawl_config,
+                    session_id="session1"
+                )
+                if result.success:
+                    print(f"Successfully crawled: {url}")
+                    await process_and_store_document(url, result.markdown_v2.raw_markdown)
+                else:
+                    print(f"Failed: {url} - Error: {result.error_message}")
+        
+        # Process all URLs in parallel with limited concurrency
+        await asyncio.gather(*[process_url(url) for url in urls])
+    finally:
+        await crawler.close()
+
+def get_pydantic_ai_docs_urls() -> List[str]:
+    """Get URLs from Pydantic AI docs sitemap."""
+    sitemap_url = "https://ai.pydantic.dev/sitemap.xml"
+    try:
+        response = requests.get(sitemap_url)
+        response.raise_for_status()
+        
+        # Parse the XML
+        root = ElementTree.fromstring(response.content)
+        
+        # Extract all URLs from the sitemap
+        namespace = {'ns': 'http://www.sitemaps.org/schemas/sitemap/0.9'}
+        urls = [loc.text for loc in root.findall('.//ns:loc', namespace)]
+        
+        return urls
+    except Exception as e:
+        print(f"Error fetching sitemap: {e}")
+        return []
+
+async def main():
+    # Get URLs from Pydantic AI docs
+    urls = get_pydantic_ai_docs_urls()
+    if not urls:
+        print("No URLs found to crawl")
+        return
+    
+    print(f"Found {len(urls)} URLs to crawl")
+    await crawl_parallel(urls)
+
+if __name__ == "__main__":
+    asyncio.run(main())

original_archon/iterations/v1-single-agent/pydantic_ai_coder.py

@@ -0,0 +1,193 @@
+from __future__ import annotations as _annotations
+
+from dataclasses import dataclass
+from dotenv import load_dotenv
+import logfire
+import asyncio
+import httpx
+import os
+
+from pydantic_ai import Agent, ModelRetry, RunContext
+from pydantic_ai.models.openai import OpenAIModel
+from openai import AsyncOpenAI
+from supabase import Client
+from typing import List
+
+load_dotenv()
+
+llm = os.getenv('LLM_MODEL', 'gpt-4o-mini')
+model = OpenAIModel(llm)
+
+logfire.configure(send_to_logfire='if-token-present')
+
+@dataclass
+class PydanticAIDeps:
+    supabase: Client
+    openai_client: AsyncOpenAI
+
+system_prompt = """
+~~ CONTEXT: ~~
+
+You are an expert at Pydantic AI - a Python AI agent framework that you have access to all the documentation to,
+including examples, an API reference, and other resources to help you build Pydantic AI agents.
+
+~~ GOAL: ~~
+
+Your only job is to help the user create an AI agent with Pydantic AI.
+The user will describe the AI agent they want to build, or if they don't, guide them towards doing so.
+You will take their requirements, and then search through the Pydantic AI documentation with the tools provided
+to find all the necessary information to create the AI agent with correct code.
+
+It's important for you to search through multiple Pydantic AI documentation pages to get all the information you need.
+Almost never stick to just one page - use RAG and the other documentation tools multiple times when you are creating
+an AI agent from scratch for the user.
+
+~~ STRUCTURE: ~~
+
+When you build an AI agent from scratch, split the agent into this files and give the code for each:
+- `agent.py`: The main agent file, which is where the Pydantic AI agent is defined.
+- `agent_tools.py`: A tools file for the agent, which is where all the tool functions are defined. Use this for more complex agents.
+- `agent_prompts.py`: A prompts file for the agent, which includes all system prompts and other prompts used by the agent. Use this when there are many prompts or large ones.
+- `.env.example`: An example `.env` file - specify each variable that the user will need to fill in and a quick comment above each one for how to do so.
+- `requirements.txt`: Don't include any versions, just the top level package names needed for the agent.
+
+~~ INSTRUCTIONS: ~~
+
+- Don't ask the user before taking an action, just do it. Always make sure you look at the documentation with the provided tools before writing any code.
+- When you first look at the documentation, always start with RAG.
+Then also always check the list of available documentation pages and retrieve the content of page(s) if it'll help.
+- Always let the user know when you didn't find the answer in the documentation or the right URL - be honest.
+- Helpful tip: when starting a new AI agent build, it's a good idea to look at the 'weather agent' in the docs as an example.
+- When starting a new AI agent build, always produce the full code for the AI agent - never tell the user to finish a tool/function.
+- When refining an existing AI agent build in a conversation, just share the code changes necessary.
+"""
+
+pydantic_ai_coder = Agent(
+    model,
+    system_prompt=system_prompt,
+    deps_type=PydanticAIDeps,
+    retries=2
+)
+
+async def get_embedding(text: str, openai_client: AsyncOpenAI) -> List[float]:
+    """Get embedding vector from OpenAI."""
+    try:
+        response = await openai_client.embeddings.create(
+            model="text-embedding-3-small",
+            input=text
+        )
+        return response.data[0].embedding
+    except Exception as e:
+        print(f"Error getting embedding: {e}")
+        return [0] * 1536  # Return zero vector on error
+
+@pydantic_ai_coder.tool
+async def retrieve_relevant_documentation(ctx: RunContext[PydanticAIDeps], user_query: str) -> str:
+    """
+    Retrieve relevant documentation chunks based on the query with RAG.
+    
+    Args:
+        ctx: The context including the Supabase client and OpenAI client
+        user_query: The user's question or query
+        
+    Returns:
+        A formatted string containing the top 5 most relevant documentation chunks
+    """
+    try:
+        # Get the embedding for the query
+        query_embedding = await get_embedding(user_query, ctx.deps.openai_client)
+        
+        # Query Supabase for relevant documents
+        result = ctx.deps.supabase.rpc(
+            'match_site_pages',
+            {
+                'query_embedding': query_embedding,
+                'match_count': 5,
+                'filter': {'source': 'pydantic_ai_docs'}
+            }
+        ).execute()
+        
+        if not result.data:
+            return "No relevant documentation found."
+            
+        # Format the results
+        formatted_chunks = []
+        for doc in result.data:
+            chunk_text = f"""
+# {doc['title']}
+
+{doc['content']}
+"""
+            formatted_chunks.append(chunk_text)
+            
+        # Join all chunks with a separator
+        return "\n\n---\n\n".join(formatted_chunks)
+        
+    except Exception as e:
+        print(f"Error retrieving documentation: {e}")
+        return f"Error retrieving documentation: {str(e)}"
+
+@pydantic_ai_coder.tool
+async def list_documentation_pages(ctx: RunContext[PydanticAIDeps]) -> List[str]:
+    """
+    Retrieve a list of all available Pydantic AI documentation pages.
+    
+    Returns:
+        List[str]: List of unique URLs for all documentation pages
+    """
+    try:
+        # Query Supabase for unique URLs where source is pydantic_ai_docs
+        result = ctx.deps.supabase.from_('site_pages') \
+            .select('url') \
+            .eq('metadata->>source', 'pydantic_ai_docs') \
+            .execute()
+        
+        if not result.data:
+            return []
+            
+        # Extract unique URLs
+        urls = sorted(set(doc['url'] for doc in result.data))
+        return urls
+        
+    except Exception as e:
+        print(f"Error retrieving documentation pages: {e}")
+        return []
+
+@pydantic_ai_coder.tool
+async def get_page_content(ctx: RunContext[PydanticAIDeps], url: str) -> str:
+    """
+    Retrieve the full content of a specific documentation page by combining all its chunks.
+    
+    Args:
+        ctx: The context including the Supabase client
+        url: The URL of the page to retrieve
+        
+    Returns:
+        str: The complete page content with all chunks combined in order
+    """
+    try:
+        # Query Supabase for all chunks of this URL, ordered by chunk_number
+        result = ctx.deps.supabase.from_('site_pages') \
+            .select('title, content, chunk_number') \
+            .eq('url', url) \
+            .eq('metadata->>source', 'pydantic_ai_docs') \
+            .order('chunk_number') \
+            .execute()
+        
+        if not result.data:
+            return f"No content found for URL: {url}"
+            
+        # Format the page with its title and all chunks
+        page_title = result.data[0]['title'].split(' - ')[0]  # Get the main title
+        formatted_content = [f"# {page_title}\n"]
+        
+        # Add each chunk's content
+        for chunk in result.data:
+            formatted_content.append(chunk['content'])
+            
+        # Join everything together
+        return "\n\n".join(formatted_content)
+        
+    except Exception as e:
+        print(f"Error retrieving page content: {e}")
+        return f"Error retrieving page content: {str(e)}"

original_archon/iterations/v1-single-agent/site_pages.sql

@@ -0,0 +1,72 @@
+-- Enable the pgvector extension
+create extension if not exists vector;
+
+-- Create the documentation chunks table
+create table site_pages (
+    id bigserial primary key,
+    url varchar not null,
+    chunk_number integer not null,
+    title varchar not null,
+    summary varchar not null,
+    content text not null,  -- Added content column
+    metadata jsonb not null default '{}'::jsonb,  -- Added metadata column
+    embedding vector(1536),  -- OpenAI embeddings are 1536 dimensions
+    created_at timestamp with time zone default timezone('utc'::text, now()) not null,
+    
+    -- Add a unique constraint to prevent duplicate chunks for the same URL
+    unique(url, chunk_number)
+);
+
+-- Create an index for better vector similarity search performance
+create index on site_pages using ivfflat (embedding vector_cosine_ops);
+
+-- Create an index on metadata for faster filtering
+create index idx_site_pages_metadata on site_pages using gin (metadata);
+
+-- Create a function to search for documentation chunks
+create function match_site_pages (
+  query_embedding vector(1536),
+  match_count int default 10,
+  filter jsonb DEFAULT '{}'::jsonb
+) returns table (
+  id bigint,
+  url varchar,
+  chunk_number integer,
+  title varchar,
+  summary varchar,
+  content text,
+  metadata jsonb,
+  similarity float
+)
+language plpgsql
+as $$
+#variable_conflict use_column
+begin
+  return query
+  select
+    id,
+    url,
+    chunk_number,
+    title,
+    summary,
+    content,
+    metadata,
+    1 - (site_pages.embedding <=> query_embedding) as similarity
+  from site_pages
+  where metadata @> filter
+  order by site_pages.embedding <=> query_embedding
+  limit match_count;
+end;
+$$;
+
+-- Everything above will work for any PostgreSQL database. The below commands are for Supabase security
+
+-- Enable RLS on the table
+alter table site_pages enable row level security;
+
+-- Create a policy that allows anyone to read
+create policy "Allow public read access"
+  on site_pages
+  for select
+  to public
+  using (true);

original_archon/iterations/v1-single-agent/streamlit_ui.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+from typing import Literal, TypedDict
+import asyncio
+import os
+
+import streamlit as st
+import json
+import logfire
+from supabase import Client
+from openai import AsyncOpenAI
+
+# Import all the message part classes
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+    UserPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    RetryPromptPart,
+    ModelMessagesTypeAdapter
+)
+from pydantic_ai_coder import pydantic_ai_coder, PydanticAIDeps
+
+# Load environment variables
+from dotenv import load_dotenv
+load_dotenv()
+
+openai_client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+supabase: Client = Client(
+    os.getenv("SUPABASE_URL"),
+    os.getenv("SUPABASE_SERVICE_KEY")
+)
+
+# Configure logfire to suppress warnings (optional)
+logfire.configure(send_to_logfire='never')
+
+class ChatMessage(TypedDict):
+    """Format of messages sent to the browser/API."""
+
+    role: Literal['user', 'model']
+    timestamp: str
+    content: str
+
+
+def display_message_part(part):
+    """
+    Display a single part of a message in the Streamlit UI.
+    Customize how you display system prompts, user prompts,
+    tool calls, tool returns, etc.
+    """
+    # system-prompt
+    if part.part_kind == 'system-prompt':
+        with st.chat_message("system"):
+            st.markdown(f"**System**: {part.content}")
+    # user-prompt
+    elif part.part_kind == 'user-prompt':
+        with st.chat_message("user"):
+            st.markdown(part.content)
+    # text
+    elif part.part_kind == 'text':
+        with st.chat_message("assistant"):
+            st.markdown(part.content)          
+
+
+async def run_agent_with_streaming(user_input: str):
+    """
+    Run the agent with streaming text for the user_input prompt,
+    while maintaining the entire conversation in `st.session_state.messages`.
+    """
+    # Prepare dependencies
+    deps = PydanticAIDeps(
+        supabase=supabase,
+        openai_client=openai_client
+    )
+
+    # Run the agent in a stream
+    async with pydantic_ai_coder.run_stream(
+        user_input,
+        deps=deps,
+        message_history= st.session_state.messages[:-1],  # pass entire conversation so far
+    ) as result:
+        # We'll gather partial text to show incrementally
+        partial_text = ""
+        message_placeholder = st.empty()
+
+        # Render partial text as it arrives
+        async for chunk in result.stream_text(delta=True):
+            partial_text += chunk
+            message_placeholder.markdown(partial_text)
+
+        # Now that the stream is finished, we have a final result.
+        # Add new messages from this run, excluding user-prompt messages
+        filtered_messages = [msg for msg in result.new_messages() 
+                            if not (hasattr(msg, 'parts') and 
+                                    any(part.part_kind == 'user-prompt' for part in msg.parts))]
+        st.session_state.messages.extend(filtered_messages)
+
+        # Add the final response to the messages
+        st.session_state.messages.append(
+            ModelResponse(parts=[TextPart(content=partial_text)])
+        )
+
+
+async def main():
+    st.title("Archon - Agent Builder")
+    st.write("Describe to me an AI agent you want to build and I'll code it for you with Pydantic AI.")
+
+    # Initialize chat history in session state if not present
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    # Display all messages from the conversation so far
+    # Each message is either a ModelRequest or ModelResponse.
+    # We iterate over their parts to decide how to display them.
+    for msg in st.session_state.messages:
+        if isinstance(msg, ModelRequest) or isinstance(msg, ModelResponse):
+            for part in msg.parts:
+                display_message_part(part)
+
+    # Chat input for the user
+    user_input = st.chat_input("What do you want to build today?")
+
+    if user_input:
+        # We append a new request to the conversation explicitly
+        st.session_state.messages.append(
+            ModelRequest(parts=[UserPromptPart(content=user_input)])
+        )
+        
+        # Display user prompt in the UI
+        with st.chat_message("user"):
+            st.markdown(user_input)
+
+        # Display the assistant's partial response while streaming
+        with st.chat_message("assistant"):
+            # Actually run the agent now, streaming the text
+            await run_agent_with_streaming(user_input)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())