Update app.py

app.py

@@ -44,6 +44,7 @@ from pinecone import Pinecone
 from PIL import Image
 from langchain_community.vectorstores import Pinecone
 from pinecone import Pinecone as PineconeClient
+from anthropic import Anthropic
 
 # Load environment variables
 load_dotenv()
@@ -69,10 +70,19 @@ vector_store = Pinecone.from_existing_index(index_name=INDEX_NAME, embedding=emb
 
 
 # Load LLM & Memory
-llm = ChatAnthropic(
-    model="claude-2", 
-    temperature=0, 
-    ANTHROPIC_API_KEY=os.getenv("ANTHROPIC_API_KEY")  
+
+class ChatAnthropicWithTokenCount(ChatAnthropic):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.anthropic_client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+
+    def count_tokens(self, text: str) -> int:
+        return self.anthropic_client.count_tokens(text)
+
+llm = ChatAnthropicWithTokenCount(
+    model="claude-2",
+    temperature=0,
+    ANTHROPIC_API_KEY=os.getenv("ANTHROPIC_API_KEY")
 )
 memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
@@ -105,7 +115,7 @@ pinecone_data = [(str(i), embedding_vectors[i], {"text": chunks[i]}) for i in ra
 vector_store.upsert(vectors=pinecone_data)
 
 # Streamlit UI
-st.set_page_config(page_title="Agentic RAG Legal Assistant", layout="wide")
+st.set_page_config(page_title="LegalAI Assistant", layout="wide")
 
 bg_image = "https://source.unsplash.com/1600x900/?law,court"
 sidebar_image = "https://source.unsplash.com/400x600/?law,justice"