Update app.py

app.py

@@ -50,40 +50,66 @@ hf_pipeline = pipeline(
 llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
 # ==========================================
-# 3. Setup RAG Chain
+# 3. Setup Dual Chains
 # ==========================================
 
-prompt_template = """
-You are a highly capable AI assistant. You possess specialized expertise in Unified Power Format (UPF) and power intent design, but you are also a helpful general conversationalist.
+# CHAIN A: The UPF Code Generator
+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.
 
-Context (Reference UPF examples, if relevant):
+Context:
 {context}
 
-User's Input:
+Power Intent Description:
 {input}
 
-Instructions:
-1. If the User's Input is asking to generate UPF code or discuss power intent, act as an expert UPF engineer. Use the Context above to generate precise, complete UPF 3.0 code. Enclose the code cleanly in ```tcl ... ``` blocks.
-2. If the User's Input is a general question, a greeting, or unrelated to UPF, respond normally as a helpful AI assistant and ignore the UPF context.
-
-Response:
+Generate only the UPF code. Do not include introductory text.
 """
-prompt = PromptTemplate.from_template(prompt_template)
-document_chain = create_stuff_documents_chain(llm, prompt)
+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)
 
+# CHAIN B: The General Conversationalist
+chat_prompt = PromptTemplate.from_template("""You are a helpful AI assistant. Answer the user's message conversationally and concisely.
+
+User: {input}
+Assistant:""")
+
 def generate_upf_code(power_intent_description):
     result = rag_chain.invoke({"input": power_intent_description})
     return result['answer'].strip()
 
+def generate_chat_response(user_message):
+    return llm.invoke(chat_prompt.format(input=user_message)).strip()
+
+
 # ==========================================
-# 4. Gradio UI
+# 4. Gradio UI & Routing Logic
 # ==========================================
 
+# Keywords that trigger the high-quality UPF generation path
+UPF_KEYWORDS = [
+    "upf", "power", "domain", "isolation", "retention", "voltage", 
+    "switch", "supply", "pst", "level shifter", "state", "intent", 
+    "create_", "set_", "connect_"
+]
+
 @spaces.GPU(duration=120) 
 def user_interaction(user_message, history):
     history = history or []
-    response = generate_upf_code(user_message)
+    
+    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:
+        # PATH A: High-quality UPF Generation
+        response = generate_upf_code(user_message)        
+        if "```" not in response:
+            response = f"```tcl\n{response}\n```"
+    else:
+        # PATH B: Normal Conversation
+        response = generate_chat_response(user_message)
 
     history.append({"role": "user", "content": user_message})
     history.append({"role": "assistant", "content": response})