agent.py

# AI Imports
from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langchain.tools import tool
from langchain.messages import HumanMessage
from langsmith import Client
from tavily import TavilyClient
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.messages import AIMessageChunk

# Data Imports
from dotenv import load_dotenv

# UX Imports
from textual.app import App, ComposeResult
from textual.widgets import Header, Footer, Input, Markdown, Label
from textual.containers import ScrollableContainer, Vertical

# Misc
import asyncio

# Load environment variables
load_dotenv()

# =================================== LLMOps =======================================
langsmith_client = Client()
tavily_client = TavilyClient()
llm = init_chat_model("gpt-4.1-mini")
checkpointer = InMemorySaver()
prompt = langsmith_client.pull_prompt("project_architecture_assistant", include_model=False)
if hasattr(prompt, 'messages') and len(prompt.messages) > 0:
    msg = prompt.messages[0]
    if hasattr(msg, 'prompt'):
       SYSTEM_PROMPT = msg.prompt.template
    elif hasattr(msg, 'content'):
       SYSTEM_PROMPT = msg.content
    else:
       SYSTEM_PROMPT = "You are a helpful assistant."
else:
    SYSTEM_PROMPT = "You are a helpful assistant."

# =================================== TOOLS ========================================

@tool("tavily_search", description="Search the web for general information based on a search query string.")
def tavily_search(query: str):
    return tavily_client.search(query, search_depth="basic", max_results=5)

@tool("tavily_extract", description="Extract raw data/content from a specific URL. The input MUST be a valid URL.")
def tavily_extract(url: str):
    result = tavily_client.extract(url, include_images=False)
    for r in result.get("results", []):
        r.pop("images", None)
    return result

tools = [tavily_search, tavily_extract]

# =================================== AGENT ========================================
agent = create_agent(model=llm, tools=tools, system_prompt=SYSTEM_PROMPT, checkpointer=checkpointer)

# =================================== APP ==========================================
class ProjectArchitect(App):
    CSS = """
    Screen { background: #1a1a1a; }
    
    #chat-container {
        height: 1fr;
        overflow-y: auto;
        padding: 1 2;
    }

    #status {
        height: 1;
        padding: 0 2;
        color: yellow;
        display: none;
    }

    #status.visible {
        display: block;
    }

    Input {
        dock: bottom;
        margin: 1 2;
        border: tall cyan;
    }

    Markdown {
        margin: 1 0;
        padding: 1 2;
        border: round #333;
    }
    """

    TITLE = "🤖 Project Architect"

    def compose(self) -> ComposeResult:
        yield Header()
        yield ScrollableContainer(id="chat-container")
        yield Label("", id="status")
        yield Input(placeholder="Ask me anything... (Ctrl+C to exit)")
        yield Footer()

    async def on_input_submitted(self, event: Input.Submitted) -> None:
        query = event.value.strip()
        if not query:
            return

        event.input.clear()
        event.input.disabled = True

        container = self.query_one("#chat-container", ScrollableContainer)
        status = self.query_one("#status", Label)

        # User message
        user_md = Markdown(f"**You:** {query}", color="green")
        await container.mount(user_md)

        # Agent response placeholder
        response_md = Markdown("")
        await container.mount(response_md)
        container.scroll_end(animate=False)

        text = ""
        config = {"configurable": {"thread_id": "1"}}
        messages = {"messages": [HumanMessage(content=query)]}

        status.add_class("visible")
        status.update("⏳ Thinking...")

        async for token, metadata in agent.astream(
            input=messages,
            config=config,
            stream_mode="messages"
        ):
            if isinstance(token, AIMessageChunk):
                if token.tool_call_chunks:
                    status.update("🔍 Researching...")
                elif token.content:
                    status.update("📝 Building analysis...")
                    text += token.content
                    await response_md.update(text)
                    container.scroll_end(animate=False)

        status.remove_class("visible")
        event.input.disabled = False
        event.input.focus()

if __name__ == "__main__":
    ProjectArchitect().run()