Update app.py

app.py

@@ -34,7 +34,6 @@ import requests
 import streamlit as st
 from fastapi import FastAPI, HTTPException
 from langchain.chains import ConversationalRetrievalChain
-from langchain.chat_models import ChatAnthropic
 from langchain.vectorstores import Pinecone
 from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 from langchain.memory import ConversationBufferMemory
@@ -45,6 +44,7 @@ from PIL import Image
 from langchain_community.vectorstores import Pinecone
 from pinecone import Pinecone as PineconeClient
 from anthropic import Anthropic
+from langchain.schema import BaseLLM
 
 # Load environment variables
 load_dotenv()
@@ -53,9 +53,6 @@ load_dotenv()
 app = FastAPI()
 
 # API Keys
-import os
-from dotenv import load_dotenv
-
 PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
 PINECONE_ENV = os.getenv("PINECONE_ENV")
 INDEX_NAME = "agenticrag"
@@ -63,29 +60,45 @@ INDEX_NAME = "agenticrag"
 if not PINECONE_API_KEY:
     raise ValueError("Pinecone API Key is missing. Please set it in environment variables.")
 
-#Initialize Hugging Face Embeddings
+# Initialize Hugging Face Embeddings
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 vector_store = Pinecone.from_existing_index(index_name=INDEX_NAME, embedding=embeddings)
 
-
-# Load LLM & Memory
-
-class ChatAnthropicWithTokenCount(ChatAnthropic):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.anthropic_client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+# Custom Anthropic LLM Wrapper
+class AnthropicLLM(BaseLLM):
+    def __init__(self, model: str, temperature: float, api_key: str):
+        super().__init__()
+        self.model = model
+        self.temperature = temperature
+        self.anthropic_client = Anthropic(api_key=api_key)
+
+    def _call(self, prompt: str, stop: list = None) -> str:
+        response = self.anthropic_client.completions.create(
+            model=self.model,
+            prompt=prompt,
+            temperature=self.temperature,
+            max_tokens_to_sample=500,  # Adjust as needed
+            stop_sequences=stop or [],
+        )
+        return response.completion
 
     def count_tokens(self, text: str) -> int:
         return self.anthropic_client.count_tokens(text)
 
-llm = ChatAnthropicWithTokenCount(
+    @property
+    def _llm_type(self) -> str:
+        return "anthropic"
+
+# Initialize Anthropic LLM
+llm = AnthropicLLM(
     model="claude-2",
     temperature=0,
-    ANTHROPIC_API_KEY=os.getenv("ANTHROPIC_API_KEY")
+    api_key=os.getenv("ANTHROPIC_API_KEY")
 )
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
+# Initialize memory
+memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
 # Build RAG Chain
 qa_chain = ConversationalRetrievalChain.from_llm(
@@ -167,4 +180,4 @@ for user_q, ai_r in st.session_state.chat_history:
     st.markdown("---")
 
 st.markdown("---")
-st.markdown("🚀 Powered by Anthropic Claude, Pinecone, and LangChain.")
+st.markdown("🚀 Powered by Anthropic Claude, Pinecone, and LangChain.")