task: [wip] set up graph and node structure

src/app.py

@@ -1,7 +1,7 @@
 import chainlit as cl
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage, SystemMessage
-from chains.design_rag import DesignRAG
+from nodes.design_rag import DesignRAG
 
 # Initialize components
 design_rag = DesignRAG()
@@ -14,8 +14,7 @@ For every user message, analyze their design preferences and requirements, consi
 3. Layout and structural needs
 4. Key visual elements
 5. Intended audience and user experience
-
-First briefly explain how you understand their requirements, then show the closest match."""
+"""
 
 @cl.on_chat_start
 async def init():
@@ -28,7 +27,7 @@ async def init():
     )
     
     # Store the LLM in the user session
-    cl.user_session.set("llm", llm)
+    cl.user_session.set("design_llm", llm)
     
     # init conversation history for each user
     cl.user_session.set("conversation_history", [
@@ -41,9 +40,9 @@ async def init():
 @cl.on_message
 async def main(message: cl.Message):
     # Get the LLM from the user session
-    llm = cl.user_session.get("llm")
-    
+    llm = cl.user_session.get("design_llm")
     conversation_history = cl.user_session.get("conversation_history")
+    
     # Add user message to history
     conversation_history.append(HumanMessage(content=message.content))
     

src/graph.py

@@ -0,0 +1,18 @@
+from typing import Annotated
+
+from typing_extensions import TypedDict
+
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import add_messages
+
+
+class State(TypedDict):
+    # Messages have the type "list". The `add_messages` function
+    # in the annotation defines how this state key should be updated
+    # (in this case, it appends messages to the list, rather than overwriting them)
+    messages: Annotated[list, add_messages]
+
+
+graph_builder = StateGraph(State)
+
+

src/nodes/design_rag.py

@@ -0,0 +1,160 @@
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain.smith import RunEvalConfig, run_on_dataset
+import os
+
+from langchain_community.vectorstores import FAISS
+from langchain.prompts import ChatPromptTemplate
+from pathlib import Path
+import json
+from typing import Dict, List, Optional
+from langchain_core.documents import Document
+from langchain.callbacks.tracers import ConsoleCallbackHandler
+
+class DesignRAG:
+    def __init__(self):
+        # Get API keys from environment
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "OPENAI_API_KEY environment variable not set. "
+                "Please set it in HuggingFace Spaces settings."
+            )
+        
+        # Initialize embedding model with explicit API key
+        self.embeddings = OpenAIEmbeddings(
+            openai_api_key=api_key
+        )
+        
+        # Load design data and create vector store
+        self.vector_store = self._create_vector_store()
+        
+        # Create retriever with tracing
+        self.retriever = self.vector_store.as_retriever(
+            search_type="similarity",
+            search_kwargs={"k": 1},
+            tags=["design_retriever"]  # Add tags for tracing
+        )
+        
+        # Create LLM with tracing
+        self.llm = ChatOpenAI(
+            temperature=0.2,
+            tags=["design_llm"]  # Add tags for tracing
+        )
+    
+    def _create_vector_store(self) -> FAISS:
+        """Create FAISS vector store from design metadata"""
+        try:
+            # Update path to look in data/designs
+            designs_dir = Path(__file__).parent.parent / "data" / "designs"
+
+            documents = []
+            
+            # 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
+                )
+            
+            print(f"Loaded {len(documents)} design documents")
+            # 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, conversation_history: List[str], num_examples: int = 1) -> str:
+        """Find similar designs based on conversation history"""
+        from langsmith import Client
+        from langchain.callbacks.tracers import ConsoleCallbackHandler
+
+        # Create LangSmith client
+        client = Client()
+        
+        # Create query generation prompt with tracing
+        query_prompt = ChatPromptTemplate.from_template(
+            """Based on this conversation history:
+            {conversation}
+            Extract the key design requirements and create a search query to find similar designs.
+            Focus on:
+            1. Visual style and aesthetics mentioned
+            2. Design categories and themes discussed
+            3. Key visual characteristics requested
+            4. Overall mood and impact desired
+            5. Any specific preferences or constraints
+            Return only the search query text, no additional explanation or analysis."""
+        ).with_config(tags=["query_generation"])
+
+        # Format conversation history
+        conversation_text = "\n".join([
+            f"{'User' if i % 2 == 0 else 'Assistant'}: {msg}"
+            for i, msg in enumerate(conversation_history)
+        ])
+        
+        # Generate optimized search query with tracing
+        query_response = await self.llm.ainvoke(
+            query_prompt.format(
+                conversation=conversation_text
+            )
+        )
+        
+        print(f"Generated query: {query_response.content}")
+        
+        # Get relevant documents with tracing
+        docs = self.retriever.get_relevant_documents(
+            query_response.content, 
+            k=num_examples,
+            callbacks=[ConsoleCallbackHandler()]  # Use standard callback instead
+        )
+        
+        # Format examples
+        examples = []
+        for doc in docs:
+            design_id = doc.metadata.get("id", "unknown")
+            content_lines = doc.page_content.strip().split("\n")
+            examples.append(
+                "\n".join(line.strip() for line in content_lines if line.strip()) +
+                f"\nURL: https://csszengarden.com/{design_id}"
+            )
+        
+        return "\n\n".join(examples)