Update app.py

app.py

@@ -45,48 +45,58 @@ hf_pipeline = pipeline(
 )
 llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
-# 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.
+# 3. Llama-3 Optimized RAG Chain
+unified_prompt_template = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
 
-Context (Reference UPF examples, if relevant):
-{context}
+You are a highly capable AI assistant specializing in Unified Power Format (UPF 3.0) and VLSI power intent design.
 
-Conversation History:
-{chat_history}
+Instructions:
+1. If the User Input asks for UPF code or power intent, act as an expert UPF engineer. Use the Context to generate precise UPF 3.0 code, enclosed in ```tcl ... ``` blocks.
+2. If the User Input is a general question or greeting, respond conversationally and concisely. Ignore the Context if it is not relevant.
 
-User Input: {input}
+Context:
+{context}<|eot_id|><|start_header_id|>user<|end_header_id|>
 
-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.
+Conversation History:
+{chat_history}
 
-Assistant:"""
+User Input: {input}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+"""
 
 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. Gradio Interface & Handling
+def get_text_content(msg_content):
+    """Safety wrapper to ensure Gradio doesn't accidentally pass JSON dicts to the LLM"""
+    if isinstance(msg_content, str):
+        return msg_content
+    elif isinstance(msg_content, list):
+        return " ".join([item.get('text', '') for item in msg_content if isinstance(item, dict) and 'text' in item])
+    return str(msg_content)
+
 def format_history(history):
     if not history:
         return "No previous conversation."
-    return "\n".join([f"{msg['role'].capitalize()}: {msg['content']}" for msg in history])
+    return "\n".join([f"{msg['role'].capitalize()}: {get_text_content(msg['content'])}" for msg in history])
 
 def user_interaction(user_message, history):
     history = history or []
+    user_text = get_text_content(user_message)
     
     response = rag_chain.invoke({
-        "input": user_message,
+        "input": user_text,
         "chat_history": format_history(history)
     })
     
-    answer = response['answer'].strip()
+    answer = response['answer'].replace("<|eot_id|>", "").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": "user", "content": user_text})
     history.append({"role": "assistant", "content": answer})
     return history, ""