task: set up chainlit langgraph

.gitignore

@@ -1,17 +1,25 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
 .env
+.venv/
+venv/
+
+# Project specific
+designs/
+.chainlit/
+dist/
+
 .env.local
 .env.development.local
 .env.test.local
 .env.production.local
 
-.venv
 .files
 
 .ipynb_checkpoints
 
-.chainlit
 .chainlit/cache
 
-/data
-
-__pycache__
+/data

src/agents/designer.py

@@ -1,19 +1,18 @@
-from typing import List, Dict
-from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages import HumanMessage
 from chains.design_rag import DesignRAG
 from .workflow import create_graph, AgentState
 from pathlib import Path
 
 class DesignerAgent:
-    def __init__(self, rag: DesignRAG):
+    def __init__(self, rag: DesignRAG) -> None:
         self.rag = rag
         self.workflow = create_graph(rag)
-        self.state: AgentState = {
-            "messages": [],
-            "html_content": self._load_default_html(),
-            "style_requirements": {},
-            "css_output": None
-        }
+        self.state: AgentState = AgentState(
+            messages=[],
+            html_content=self._load_default_html(),
+            style_requirements={},
+            css_output=None
+        )
     
     def _load_default_html(self) -> str:
         """Load default CSS Zen Garden HTML"""
@@ -28,14 +27,11 @@ class DesignerAgent:
     
     async def process(self, message: str) -> str:
         """Process a message through the workflow"""
-        # Add message to state
-        self.state["messages"].append(HumanMessage(content=message))
-        
-        # Run workflow
-        next_state = await self.workflow.invoke(self.state)
-        
-        # Update state
-        self.state = next_state
-        
-        # Return last response
-        return self.state["messages"][-1].content 
+        try:
+            self.state["messages"].append(HumanMessage(content=message))
+            next_state = await self.workflow.ainvoke(self.state)
+            self.state = next_state
+            return self.state["messages"][-1].content
+        except Exception as e:
+            print(f"Error in process: {str(e)}")
+            return "I encountered an error processing your message." 

src/agents/workflow.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Annotated, TypedDict
+from typing import Dict, List, Annotated, TypedDict, Callable, Any
 from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.graph import StateGraph, END
 from langchain_core.messages import BaseMessage
@@ -6,6 +6,7 @@ from chains.design_rag import DesignRAG
 from langchain.prompts import ChatPromptTemplate
 import chainlit as cl
 import json
+from langchain_openai import ChatOpenAI
 
 # Define state types
 class AgentState(TypedDict):
@@ -14,35 +15,47 @@ class AgentState(TypedDict):
     style_requirements: Dict
     css_output: str | None
 
-# Node functions
-async def conversation_node(state: AgentState, rag: DesignRAG):
+# At the top level
+llm = ChatOpenAI(temperature=0.2)
+
+def should_end(state: AgentState) -> str:
+    """Determine if the conversation should end"""
+    if state["messages"] and isinstance(state["messages"][-1], HumanMessage):
+        if "!done" in state["messages"][-1].content.lower():
+            return "end"
+    return "conversation"
+
+async def conversation_node(state: AgentState) -> AgentState:
     """Handle conversation and requirement gathering"""
-    # Get last message
     last_message = state["messages"][-1]
+    print("last message received")
+    print(last_message)
     
     if not isinstance(last_message, HumanMessage):
-        return {"messages": state["messages"]}
+        return state
     
-    # Check for style requirements readiness
     if "!generate" in last_message.content.lower():
-        # Extract style requirements from conversation
-        requirements = await extract_requirements(state["messages"], rag)
-        return {
-            "messages": state["messages"],
-            "style_requirements": requirements,
-            "next": "generate_css"
-        }
+        requirements = await extract_requirements(state["messages"])
+        await cl.Message("I'll generate CSS based on your requirements...").send()
+        return AgentState(
+            messages=state["messages"],
+            html_content=state["html_content"],
+            style_requirements=requirements,
+            css_output=state["css_output"]
+        )
     
-    # Normal conversation - get context and respond
-    response = await rag.query(last_message.content)
-    state["messages"].append(AIMessage(content=response))
+    # Normal conversation - just acknowledge and guide
+    response = "I understand. Tell me more about the design you're looking for, or type !generate when you're ready to create the CSS."
+    await cl.Message(content=response).send()
     
-    return {"messages": state["messages"]}
+    return AgentState(
+        messages=[*state["messages"], AIMessage(content=response)],
+        html_content=state["html_content"],
+        style_requirements=state.get("style_requirements", {}),
+        css_output=state.get("css_output")
+    )
 
-async def extract_requirements(
-    messages: List[BaseMessage], 
-    rag: DesignRAG
-) -> Dict:
+async def extract_requirements(messages: List[BaseMessage]) -> Dict:
     """Extract style requirements from conversation history"""
     # Combine messages into context
     context = "\n".join([
@@ -65,8 +78,8 @@ async def extract_requirements(
     }}
     """
     
-    # Get requirements through RAG system
-    response = await rag.llm.ainvoke(prompt)
+    # Get requirements through LLM
+    response = await llm.ainvoke(prompt)
     return json.loads(response.content)
 
 async def generate_css_node(state: AgentState, rag: DesignRAG) -> AgentState:
@@ -107,37 +120,44 @@ async def generate_css_node(state: AgentState, rag: DesignRAG) -> AgentState:
         )
     )
     
-    # Store generated CSS
-    state["css_output"] = response.content
-    
-    # Add completion message
-    state["messages"].append(AIMessage(content="""I've generated the CSS based on your requirements. 
+    css_message = f"""I've generated the CSS based on your requirements. 
     Here's what I created:
     
     ```css
-    {css}
+    {response.content}
     ```
     
     Would you like me to explain any part of the design or make any adjustments?
-    """.format(css=response.content)))
+    Type !done when you're satisfied with the result."""
+    
+    await cl.Message(content=css_message).send()
     
-    return state
+    return AgentState(
+        messages=[*state["messages"], AIMessage(content=css_message)],
+        html_content=state["html_content"],
+        style_requirements=state["style_requirements"],
+        css_output=response.content
+    )
 
 def create_graph(rag: DesignRAG) -> StateGraph:
     """Create the workflow graph"""
-    # Create graph
     workflow = StateGraph(AgentState)
     
-    # Add nodes
-    workflow.add_node("conversation", lambda s: conversation_node(s, rag))
+    # Add nodes directly
+    workflow.add_node("conversation", conversation_node)
     workflow.add_node("generate_css", lambda s: generate_css_node(s, rag))
     
     # Add edges
-    workflow.add_edge("conversation", "conversation")
-    workflow.add_edge("conversation", "generate_css")
-    workflow.add_edge("generate_css", END)
+    workflow.add_conditional_edges(
+        "conversation",
+        should_end,  # Use our existing should_end function
+        {
+            "conversation": "conversation",
+            "generate_css": "generate_css",
+            "end": END
+        }
+    )
+    workflow.add_edge("generate_css", "conversation")
     
-    # Set entry point
     workflow.set_entry_point("conversation")
-    
     return workflow.compile() 

src/app.py

@@ -1,23 +1,30 @@
 import chainlit as cl
-from langchain_core.messages import HumanMessage, AIMessage
 from chains.design_rag import DesignRAG
 from agents.designer import DesignerAgent
 
+# Initialize these once at module level
+design_rag = DesignRAG()
+designer = DesignerAgent(rag=design_rag)
+
 @cl.on_chat_start
-async def start():
+async def start() -> None:
     """Initialize the chat session"""
-    # Initialize RAG system
-    design_rag = DesignRAG()
-    # Initialize designer agent
-    designer = DesignerAgent(rag=design_rag)
-    
     # Store in user session
     cl.user_session.set("designer", designer)
+    
+    # Send welcome message
+    await cl.Message(
+        content="Welcome! I'm here to help you imagine a unique design. What style are you looking for?"
+    ).send()
 
 @cl.on_message
-async def main(message: cl.Message):
+async def main(message: cl.Message) -> None:
     """Handle incoming messages"""
     designer = cl.user_session.get("designer")
+    if designer is None:
+        # Reinitialize if missing
+        designer = DesignerAgent(rag=design_rag)
+        cl.user_session.set("designer", designer)
     
     # Process message through designer agent
     response = await designer.process(message.content)

src/chains/design_rag.py

@@ -5,7 +5,8 @@ from langchain_community.vectorstores import FAISS
 from langchain.prompts import ChatPromptTemplate
 from pathlib import Path
 import json
-from typing import Dict
+from typing import Dict, List, Optional
+from langchain_core.documents import Document
 
 class DesignRAG:
     def __init__(self):
@@ -18,78 +19,68 @@ class DesignRAG:
         # Create retriever
         self.retriever = self.vector_store.as_retriever(
             search_type="similarity",
-            search_kwargs={"k": 3}
+            search_kwargs={"k": 5}
         )
         
         # Create LLM
-        self.llm = ChatOpenAI(temperature=0.7)
-        
-        # Create the RAG chain
-        self.chain = self._create_chain()
+        self.llm = ChatOpenAI(temperature=0.2)
     
-    def _create_vector_store(self):
+    def _create_vector_store(self) -> FAISS:
         """Create FAISS vector store from design metadata"""
-        designs_dir = Path("designs")
-        documents = []
-        
-        # Load all metadata files
-        for design_dir in designs_dir.glob("**/metadata.json"):
-            with open(design_dir, "r") as f:
-                metadata = json.load(f)
-                
-            # Create document text from metadata
-            text = f"""
-            Design {metadata['id']}:
-            Description: {metadata.get('description', '')}
-            Categories: {', '.join(metadata.get('categories', []))}
-            Visual Characteristics: {', '.join(metadata.get('visual_characteristics', []))}
-            """
+        try:
+            # Update path to look in data/designs
+            designs_dir = Path("data/designs")
             
-            # Load associated CSS
-            css_path = design_dir.parent / "style.css"
-            if css_path.exists():
-                with open(css_path, "r") as f:
-                    css = f.read()
-                text += f"\nCSS:\n{css}"
+            documents = []
             
-            documents.append({
-                "page_content": text,
-                "metadata": {
-                    "id": metadata["id"],
-                    "url": metadata.get("url", ""),
-                }
-            })
-        
-        # Create and return vector store
-        return FAISS.from_documents(documents, self.embeddings)
-    
-    def _create_chain(self):
-        """Create the RAG processing chain"""
-        # Define prompt template
-        template = """You are a design assistant helping to find and adapt CSS Zen Garden designs.
-        Use the following similar designs to inform your response:
-        
-        {context}
-        
-        Based on these examples, help the user with their request:
-        {question}
-        """
-        
-        prompt = ChatPromptTemplate.from_template(template)
-        
-        # Create and return chain
-        chain = (
-            {"context": self.retriever, "question": RunnablePassthrough()}
-            | prompt
-            | self.llm
-            | StrOutputParser()
-        )
-        
-        return chain
-    
-    async def query(self, question: str) -> str:
-        """Process a query through the RAG system"""
-        return await self.chain.ainvoke(question)
+            # Load all metadata files
+            for design_dir in designs_dir.glob("**/metadata.json"):
+                try:
+                    with open(design_dir, "r") as f:
+                        metadata = json.load(f)
+                    
+                    # Create document text from metadata with safe gets
+                    text = f"""
+                    Design {metadata.get('id', 'unknown')}:
+                    Description: {metadata.get('description', 'No description available')}
+                    Categories: {', '.join(metadata.get('categories', []))}
+                    Visual Characteristics: {', '.join(metadata.get('visual_characteristics', []))}
+                    """
+                    
+                    # Load associated CSS
+                    css_path = design_dir.parent / "style.css"
+                    if css_path.exists():
+                        with open(css_path, "r") as f:
+                            css = f.read()
+                        text += f"\nCSS:\n{css}"
+                    
+                    # Create Document object with minimal metadata
+                    documents.append(
+                        Document(
+                            page_content=text.strip(),
+                            metadata={
+                                "id": metadata.get('id', 'unknown'),
+                                "path": str(design_dir.parent)
+                            }
+                        )
+                    )
+                except Exception as e:
+                    print(f"Error processing design {design_dir}: {e}")
+                    continue
+            
+            if not documents:
+                print("Warning: No valid design documents found")
+                # Create empty vector store with a placeholder document
+                return FAISS.from_documents(
+                    [Document(page_content="No designs available", metadata={"id": "placeholder"})],
+                    self.embeddings
+                )
+            
+            # Create and return vector store
+            return FAISS.from_documents(documents, self.embeddings)
+        except Exception as e:
+            print(f"Error creating vector store: {str(e)}")
+            raise
     
     async def query_similar_designs(self, requirements: Dict) -> str:
         """Find similar designs based on requirements"""
@@ -103,7 +94,7 @@ class DesignRAG:
         """
         
         # Get similar documents
-        docs = self.retriever.get_relevant_documents(query)
+        docs = await self.retriever.get_relevant_documents(query)
         
         # Format examples
         examples = []