Update app.py

app.py

@@ -7,17 +7,17 @@ from langchain_classic.chains import create_retrieval_chain
 from langchain_classic.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import PromptTemplate
 
-# 1. Load Vector Store
+# 1. Vector Store Setup
 embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
 vectorstore = FAISS.load_local(
     "faiss_upf_index", 
     embeddings, 
     allow_dangerous_deserialization=True
 )
-retriever = vectorstore.as_retriever()
+retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
 
-# 2. Load Model & Pipeline in 4-bit
-model_id = "anirudh248/llama3-upf-generator"
+# 2. Model & Pipeline Initialization
+model_id = "anirudh248/upf-code-generator"
 
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -45,59 +45,53 @@ hf_pipeline = pipeline(
 )
 llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
-# 3. Setup RAG Chain & Chat Prompt
-upf_prompt_template = """
-You are an expert in Unified Power Format (UPF). Generate a precise and complete UPF code block based on the following power intent. Use the retrieved context as a reference. The code must be correct and adhere to UPF 3.0 standards.
+# 3. Unified RAG Chain
+unified_prompt_template = """You are a highly capable AI assistant specializing in Unified Power Format (UPF 3.0) and VLSI power intent design.
 
-Context:
+Context (Reference UPF examples, if relevant):
 {context}
 
-Power Intent Description:
-{input}
+Conversation History:
+{chat_history}
 
-Generate only the UPF code. Do not include introductory text.
-"""
-upf_prompt = PromptTemplate.from_template(upf_prompt_template)
-document_chain = create_stuff_documents_chain(llm, upf_prompt)
-rag_chain = create_retrieval_chain(retriever, document_chain)
+User Input: {input}
 
-chat_prompt = PromptTemplate.from_template(
-    "You are a helpful AI assistant. Answer the user's message conversationally and concisely.\n\nUser: {input}\nAssistant:"
-)
+Instructions:
+1. Analyze the User Input. If it requests UPF code, power intent design, or hardware specifics, act as an expert UPF engineer. Use the Context to generate precise UPF 3.0 code, enclosed completely in ```tcl ... ``` blocks.
+2. If the User Input is a general question, greeting, or unrelated to UPF, respond normally as a helpful AI assistant. Ignore the UPF Context and do not generate code.
 
-def generate_upf_code(power_intent_description):
-    result = rag_chain.invoke({"input": power_intent_description})
-    return result['answer'].strip()
+Assistant:"""
 
-def generate_chat_response(user_message):
-    return llm.invoke(chat_prompt.format(input=user_message)).strip()
+unified_prompt = PromptTemplate.from_template(unified_prompt_template)
+document_chain = create_stuff_documents_chain(llm, unified_prompt)
+rag_chain = create_retrieval_chain(retriever, document_chain)
 
-# 4. Routing Logic & Gradio Interface
-UPF_KEYWORDS = [
-    "upf", "power", "domain", "isolation", "retention", "voltage", 
-    "switch", "supply", "pst", "level shifter", "state", "intent", 
-    "create_", "set_", "connect_"
-]
+# 4. Gradio Interface & Handling
+def format_history(history):
+    if not history:
+        return "No previous conversation."
+    return "\n".join([f"{msg['role'].capitalize()}: {msg['content']}" for msg in history])
 
 def user_interaction(user_message, history):
     history = history or []
-    msg_lower = user_message.lower()
     
-    is_upf_request = any(kw in msg_lower for kw in UPF_KEYWORDS) or len(msg_lower.split()) > 12
-
-    if is_upf_request:
-        response = generate_upf_code(user_message)
-        if "```" not in response:
-            response = f"```tcl\n{response}\n```"
-    else:
-        response = generate_chat_response(user_message)
+    response = rag_chain.invoke({
+        "input": user_message,
+        "chat_history": format_history(history)
+    })
+    
+    answer = response['answer'].strip()
+    
+    upf_syntax_hints = ["create_power_domain", "set_isolation", "set_retention", "create_supply_port", "create_power_switch"]
+    if "```" not in answer and any(kw in answer.lower() for kw in upf_syntax_hints):
+        answer = f"```tcl\n{answer}\n```"
 
     history.append({"role": "user", "content": user_message})
-    history.append({"role": "assistant", "content": response})
+    history.append({"role": "assistant", "content": answer})
     return history, ""
 
 with gr.Blocks() as interface:
-    gr.Markdown("# UPF Code Generator with Llama 3 & RAG")
+    gr.Markdown("# ⚡ UPF Code Generator with Llama 3 & RAG")
     
     chatbot = gr.Chatbot(label="Chat History", elem_id="chatbot")