Update main.py

main.py

@@ -1,26 +1,18 @@
 import os
 import streamlit as st
-from dotenv import load_dotenv
 from pydantic import BaseModel
-from langchain_community.llms import HuggingFaceHub
+from langchain_anthropic import ChatAnthropic
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import PydanticOutputParser
 from langchain.agents import create_tool_calling_agent, AgentExecutor
 from tools import search_tool, wiki_tool, save_tool
 
-# ✅ Get the token from environment variables
-hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+# ✅ Get key from Hugging Face Secrets
+claude_key = os.getenv("CLAUDE_API_KEY")
 
-# ✅ Check and raise an error if it's missing
-if hf_token is None:
-    raise ValueError("Missing Hugging Face API token. Please set the 'HUGGINGFACEHUB_API_TOKEN' secret in your HF space settings.")
+if claude_key is None:
+    raise ValueError("Claude API key not found. Please set CLAUDE_API_KEY in Hugging Face Secrets.")
 
-# ✅ Use HuggingFace Mistral model (free), Initialize your LLM using the token
-llm = HuggingFaceHub(
-    repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-    huggingfacehub_api_token=hf_token,
-    model_kwargs={"temperature": 0.7, "max_new_tokens": 500},
-)
 
 # Define the response format
 class ResearchResponse(BaseModel):
@@ -29,6 +21,9 @@ class ResearchResponse(BaseModel):
     sources: list[str]
     tools_used: list[str]
 
+# Initialize Claude model with low token usage
+llm = ChatAnthropic(model="claude-3-5-sonnet-20241022", max_tokens=300, anthropic_api_key=claude_key)
+
 # Create parser
 parser = PydanticOutputParser(pydantic_object=ResearchResponse)
 
@@ -38,7 +33,7 @@ prompt = ChatPromptTemplate.from_messages(
         (
             "system",
             """
-            You are an AI assistant that will help with any general questions.
+            You are an AI assistant that will help of any general questions.
             Answer the user query and use necessary tools. Don't answer any offensive questions or any swearings.
             Wrap the output in this format and provide no other text\n{format_instructions}
             """,
@@ -60,19 +55,24 @@ agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False)
 st.set_page_config(page_title="TechTales AI - Your own AI Powered Assistant", layout="centered")
 st.title("💡 TechTales AI - Your own AI Powered Assistant")
 
+# Initialize chat history in session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+# Display chat history
 for msg in st.session_state.messages:
     st.chat_message(msg["role"]).write(msg["content"])
 
+# Chat input with Enter button
 query = st.text_input("I am your personal Assistant, Ask me anything...", key="query_input")
 submit_button = st.button("Enter")
 
 if submit_button and query:
+    # Add user message to history
     st.session_state.messages.append({"role": "user", "content": query})
     st.chat_message("user").write(query)
 
+    # Get AI response
     raw_response = agent_executor.invoke({"query": query})
 
     try:
@@ -81,12 +81,15 @@ if submit_button and query:
     except Exception:
         response_text = "I'm sorry, I couldn't process that request."
 
+    # Add AI response to history
     st.session_state.messages.append({"role": "assistant", "content": response_text})
     st.chat_message("assistant").write(response_text)
 
+# Clear Chat Button
 if st.button("🗑️ Clear Chat"):
     st.session_state.messages = []
-    st.rerun()
+    st.rerun()  # Corrected method
 
+# Developer Name Display
 st.markdown("---")
 st.markdown("**Developed by: Pankaj Kumar**", unsafe_allow_html=True)