├── .DS_Store ├── .gitignore ├── LICENSE ├── README.md ├── app.py ├── requirements.txt ├── streamlit_app.py ├── utils.py └── vector_db.py /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/alonlavian/RAGent/77b28c388cb194a21c71611d94ab6b330f779ae0/.DS_Store -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Python 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | *.so 6 | .Python 7 | env/ 8 | build/ 9 | develop-eggs/ 10 | dist/ 11 | downloads/ 12 | eggs/ 13 | .eggs/ 14 | lib/ 15 | lib64/ 16 | parts/ 17 | sdist/ 18 | var/ 19 | wheels/ 20 | *.egg-info/ 21 | .installed.cfg 22 | *.egg 23 | 24 | # Virtual Environment 25 | venv/ 26 | ENV/ 27 | 28 | # IDE 29 | .idea/ 30 | .vscode/ 31 | 32 | # Project specific 33 | chroma_db/ 34 | *.pdf 35 | .env -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Alon Lavian 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # The King RAGent 2 | 3 | The King RAGent is an AI-powered research assistant that leverages vector databases and external APIs to provide comprehensive answers to user queries. This application is built using Streamlit for the frontend and integrates various backend services for document processing and information retrieval. 4 | 5 | ## Features 6 | 7 | - **PDF Upload and Processing**: Upload PDF documents to initialize the vector store. 8 | - **AI-Powered Query Handling**: Use AI models to process and respond to user queries. 9 | - **Web Search Integration**: Augment responses with web search results for enhanced accuracy. 10 | - **Dry Run Mode**: Test the application without making actual API calls or database operations. 11 | 12 | ## Installation 13 | 14 | 1. Clone the repository: 15 | ```bash 16 | git clone https://github.com/yourusername/ragents.git 17 | cd ragents 18 | ``` 19 | 20 | 2. Install the required packages: 21 | ```bash 22 | pip install -r requirements.txt 23 | ``` 24 | 25 | 3. Set up environment variables: 26 | - Create a `.env` file in the root directory. 27 | - Add your API keys and other necessary configurations. 28 | 29 | ## Usage 30 | 31 | 1. Run the Streamlit app: 32 | ```bash 33 | streamlit run streamlit_app.py 34 | ``` 35 | 36 | 2. Open your browser and navigate to the local URL provided by Streamlit. 37 | 38 | ## Dry Run Mode 39 | 40 | The application now supports a `dry_run` mode, which allows you to test the app without making actual API calls or database operations. This is useful for development and testing purposes. 41 | 42 | ### How to Enable Dry Run Mode 43 | 44 | - **Streamlit Interface**: Use the "🔧 Dry Run Mode" checkbox in the sidebar to toggle dry run mode on or off. 45 | - **Backend Logic**: The `dry_run` parameter is propagated through the application, affecting the following components: 46 | - **VectorStore**: Skips database loading and returns mock data. 47 | - **call_claude_rag**: Returns a mock response instead of calling the Claude API. 48 | - **assess_confidence**: Returns a mock confidence level. 49 | - **synthesize_information**: Returns a mock synthesis of information. 50 | - **call_tavily_web_search**: Returns mock search results. 51 | 52 | ## Contributing 53 | 54 | Contributions are welcome! Please fork the repository and submit a pull request for any improvements or bug fixes. 55 | 56 | ## License 57 | 58 | This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details. -------------------------------------------------------------------------------- /app.py: -------------------------------------------------------------------------------- 1 | import time 2 | from vector_db import VectorStore 3 | from utils import call_claude_rag, call_tavily_web_search, assess_confidence, synthesize_information 4 | import logging 5 | 6 | def main(user_query, initialize=False, pdf_path=None, dry_run=False): # Added dry_run parameter 7 | logging.debug(f"[main] Starting process with query: {user_query[:100]}...") 8 | if dry_run: 9 | logging.info("[main] Running in dry run mode") 10 | 11 | vector_store = VectorStore(persist_directory="./chroma_db") 12 | 13 | if initialize and pdf_path: 14 | logging.info(f"[main] Initializing with PDF: {pdf_path}") 15 | vector_store.initialize_from_pdf(pdf_path) 16 | return "PDF processed and ready for questions!" 17 | elif not vector_store.db and not initialize: 18 | logging.error("[main] No vector store found") 19 | return "Error: Vector store not initialized. Please provide a PDF file first." 20 | elif initialize: 21 | logging.error("[main] PDF path not provided for initialization") 22 | return "Please provide a PDF file to initialize the system." 23 | 24 | logging.debug("[main] Starting query processing") 25 | internal_response = vector_store.search(user_query) 26 | 27 | logging.debug("[main] Getting RAG response") 28 | rag_response = call_claude_rag(user_query, internal_response, dry_run) 29 | # time.sleep(1) 30 | 31 | if isinstance(rag_response, list): 32 | rag_response = " ".join(str(item) for item in rag_response) 33 | 34 | logging.debug("[main] Assessing confidence") 35 | if assess_confidence(rag_response, dry_run): 36 | logging.info("[main] Returning high-confidence response") 37 | return rag_response 38 | 39 | logging.debug("[main] Confidence low, performing web search") 40 | web_results = call_tavily_web_search(user_query, dry_run) 41 | time.sleep(1) 42 | 43 | logging.debug("[main] Synthesizing information") 44 | synthesized_response = synthesize_information(rag_response, web_results, dry_run) 45 | 46 | logging.info("[main] Returning synthesized response") 47 | return synthesized_response -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | requests==2.31.0 2 | anthropic 3 | langchain 4 | langchain-community 5 | langchain-core 6 | langchain-chroma 7 | chromadb 8 | sentence-transformers 9 | pypdf==3.17.4 10 | pydantic>=2.0.0,<3.0.0 11 | pydantic-core>=2.0.0,<3.0.0 12 | tavily-python 13 | langchain-huggingface 14 | python-dotenv 15 | streamlit -------------------------------------------------------------------------------- /streamlit_app.py: -------------------------------------------------------------------------------- 1 | import streamlit as st 2 | from app import main 3 | import os 4 | 5 | st.set_page_config(page_title="The King RAGent", page_icon="👑", layout="wide") 6 | 7 | def initialize_session_state(): 8 | if "messages" not in st.session_state: 9 | st.session_state.messages = [] 10 | if "initialized" not in st.session_state: 11 | st.session_state.initialized = False 12 | 13 | def display_chat_history(): 14 | for message in st.session_state.messages: 15 | with st.chat_message(message["role"]): 16 | st.markdown(message["content"]) 17 | 18 | def main_chat(): 19 | # Create two columns: sidebar and main content 20 | with st.sidebar: 21 | st.title("Settings") 22 | dry_run = st.checkbox("🔧 Dry Run Mode", value=False, help="Run without making API calls") 23 | if dry_run: 24 | st.warning("Test mode active - using stub responses") 25 | 26 | # Main content 27 | st.title("👑 The King RAGent") 28 | st.markdown("### Your AI Research Assistant") 29 | 30 | initialize_session_state() 31 | 32 | # PDF Upload 33 | if not st.session_state.initialized: 34 | uploaded_file = st.file_uploader("Upload a PDF file", type="pdf") 35 | if uploaded_file is not None: 36 | # Save the uploaded file temporarily 37 | with open("temp.pdf", "wb") as f: 38 | f.write(uploaded_file.getvalue()) 39 | 40 | # Initialize the vector store 41 | response = main(None, initialize=True, pdf_path="temp.pdf", dry_run=dry_run) 42 | st.session_state.initialized = True 43 | os.remove("temp.pdf") # Clean up 44 | st.rerun() 45 | 46 | # Display chat history 47 | display_chat_history() 48 | 49 | # Chat input 50 | if prompt := st.chat_input("What would you like to know?"): 51 | # Add user message to chat history 52 | st.session_state.messages.append({"role": "user", "content": prompt}) 53 | 54 | # Display user message 55 | with st.chat_message("user"): 56 | st.markdown(prompt) 57 | 58 | # Get AI response 59 | with st.chat_message("assistant"): 60 | with st.spinner("Thinking..."): 61 | response = main(prompt, initialize=False, dry_run=dry_run) 62 | 63 | # Add AI response to chat history 64 | st.session_state.messages.append({"role": "assistant", "content": response}) 65 | 66 | # Rerun to update the chat display 67 | st.rerun() 68 | 69 | if __name__ == "__main__": 70 | main_chat() -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import time 2 | from anthropic import Anthropic 3 | import logging 4 | import requests 5 | from tavily import TavilyClient 6 | from dotenv import load_dotenv 7 | import os 8 | 9 | # Load environment variables 10 | load_dotenv() 11 | 12 | # Initialize the Anthropic client 13 | client = Anthropic(api_key=os.getenv('CLAUDE_API_KEY')) 14 | 15 | def call_claude_rag(query, context, dry_run=False): 16 | if dry_run: 17 | logging.info("[call_claude_rag] Dry run mode - returning mock response") 18 | return "Mock response: This is a dry run response." 19 | 20 | logging.debug(f"[call_claude_rag] Processing query: {query[:100]}...") # Log first 100 chars of query 21 | system_prompt = f"""You are a helpful assistant that answers questions based on the provided context. 22 | If the context doesn't contain enough information to answer confidently, indicate that. 23 | 24 | Context: 25 | {context} 26 | """ 27 | 28 | try: 29 | message = client.messages.create( 30 | model="claude-3-5-sonnet-20241022", 31 | max_tokens=512, 32 | system=system_prompt, 33 | messages=[ 34 | {"role": "user", "content": query} 35 | ] 36 | ) 37 | logging.debug("[call_claude_rag] Successfully received response from Claude") 38 | return message.content 39 | except Exception as e: 40 | logging.error(f"[call_claude_rag] Error calling Claude API: {e}") 41 | return "" 42 | 43 | def call_tavily_web_search(query): 44 | logging.debug(f"[call_tavily_web_search] Searching for: {query}") 45 | try: 46 | tavily_client = TavilyClient(api_key=os.getenv('TAVILY_API_KEY')) 47 | context = tavily_client.get_search_context(query=query) 48 | logging.debug("[call_tavily_web_search] Successfully received search results") 49 | return context 50 | except Exception as e: 51 | logging.error(f"[call_tavily_web_search] Error in web search: {e}") 52 | return "" 53 | 54 | def assess_confidence(response, dry_run=False): 55 | if dry_run: 56 | logging.info("[assess_confidence] Dry run mode - returning mock confidence") 57 | return True # or False, depending on what you want to simulate 58 | 59 | """ 60 | Use Claude to evaluate the confidence of the response. 61 | Returns True if confident, False if not confident. 62 | """ 63 | logging.debug("[assess_confidence] Evaluating response confidence") 64 | system_prompt = """You are a confidence assessment expert. Evaluate the confidence level of the following response. 65 | Consider: 66 | 1. Completeness of the answer 67 | 2. Specificity and precision 68 | 3. Presence of uncertainty markers (e.g., "might", "maybe", "I'm not sure") 69 | 4. Consistency of information 70 | 71 | You must respond in valid JSON format with exactly these fields: 72 | { 73 | "confidence_score": , 74 | "reasoning": "" 75 | } 76 | """ 77 | 78 | try: 79 | message = client.messages.create( 80 | model="claude-3-5-sonnet-20241022", 81 | max_tokens=512, 82 | system=system_prompt, 83 | messages=[ 84 | {"role": "user", "content": f"Response to evaluate: {response}"} 85 | ] 86 | ) 87 | 88 | evaluation = message.content 89 | logging.debug(f"[assess_confidence] Raw evaluation: {evaluation}") 90 | 91 | # Handle TextBlock wrapper if present 92 | if isinstance(evaluation, list): 93 | evaluation = evaluation[0].text if hasattr(evaluation[0], 'text') else str(evaluation[0]) 94 | if 'TextBlock(text=' in str(evaluation): 95 | evaluation = str(evaluation).split('text=\'')[1].split('\',')[0] 96 | 97 | logging.debug(f"[assess_confidence] Cleaned evaluation: {evaluation}") 98 | 99 | try: 100 | # Try parsing as JSON 101 | import json 102 | parsed = json.loads(evaluation) 103 | confidence_score = float(parsed["confidence_score"]) 104 | logging.debug(f"[assess_confidence] Parsed confidence score: {confidence_score}") 105 | return confidence_score > 0.7 106 | except (json.JSONDecodeError, KeyError) as e: 107 | logging.warning(f"[assess_confidence] JSON parsing failed: {e}") 108 | 109 | # Fallback to string parsing 110 | if '"confidence_score":' in evaluation: 111 | score_text = evaluation.split('"confidence_score":')[1].split(',')[0].strip() 112 | confidence_score = float(score_text) 113 | logging.debug(f"[assess_confidence] Extracted confidence score: {confidence_score}") 114 | return confidence_score > 0.7 115 | 116 | logging.warning("[assess_confidence] Could not extract confidence score from response") 117 | logging.debug(f"[assess_confidence] Response content: {evaluation}") 118 | return False 119 | 120 | except Exception as e: 121 | logging.error(f"[assess_confidence] Error in confidence assessment: {e}") 122 | return False 123 | 124 | def synthesize_information(internal_response, web_results, dry_run=False): 125 | if dry_run: 126 | logging.info("[synthesize_information] Dry run mode - returning mock synthesis") 127 | return "Mock synthesis: This is a dry run synthesis." 128 | 129 | logging.debug("[synthesize_information] Starting synthesis of information") 130 | system_prompt = """Synthesize the internal knowledge and web search results into a comprehensive response. 131 | Format your response in a clear, human-readable way: 132 | 1. Use bullet points for lists 133 | 2. Add line breaks between sections 134 | 3. Bold important numbers and dates 135 | 4. Present the information in a conversational tone 136 | 5. Highlight key findings at the beginning 137 | 6. If there are discrepancies between sources, explain them clearly 138 | 139 | Structure your response with these sections: 140 | - Key Finding 141 | - Details from Internal Knowledge 142 | - Details from Web Search 143 | - Additional Context (if any) 144 | """ 145 | 146 | try: 147 | message = client.messages.create( 148 | model="claude-3-5-sonnet-20241022", 149 | max_tokens=512, 150 | system=system_prompt, 151 | messages=[ 152 | {"role": "user", "content": f"Internal Knowledge: {internal_response}\n\nWeb Results: {web_results}"} 153 | ] 154 | ) 155 | 156 | response = message.content 157 | if isinstance(response, list): 158 | response = response[0].text if hasattr(response[0], 'text') else str(response[0]) 159 | 160 | if response.startswith('TextBlock('): 161 | response = response.split('text=\'', 1)[1].rsplit('\')', 1)[0] 162 | 163 | logging.debug("[synthesize_information] Successfully synthesized information") 164 | return response 165 | except Exception as e: 166 | logging.error(f"[synthesize_information] Error in synthesis: {e}") 167 | return "Error synthesizing information" -------------------------------------------------------------------------------- /vector_db.py: -------------------------------------------------------------------------------- 1 | # Set environment variable before imports 2 | import os 3 | os.environ["TOKENIZERS_PARALLELISM"] = "false" 4 | 5 | import warnings 6 | warnings.filterwarnings('ignore', category=DeprecationWarning) 7 | 8 | import logging 9 | from langchain_chroma import Chroma 10 | from langchain_community.document_loaders import PyPDFLoader 11 | from langchain_text_splitters import CharacterTextSplitter 12 | from langchain_community.embeddings import HuggingFaceEmbeddings 13 | import chromadb 14 | 15 | # Set up logging 16 | logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') 17 | 18 | class VectorStore: 19 | def __init__(self, model_name="all-MiniLM-L6-v2", persist_directory="./chroma_db", dry_run=False): 20 | """Initialize the vector store with the specified embedding model.""" 21 | self.dry_run = dry_run 22 | logging.debug(f"[__init__] Initializing VectorStore with model: {model_name}, dry_run: {self.dry_run}") 23 | self.embedding_function = HuggingFaceEmbeddings(model_name=model_name) 24 | self.persist_directory = persist_directory 25 | 26 | if not self.dry_run: 27 | # Try to load existing DB 28 | if os.path.exists(persist_directory): 29 | try: 30 | self.db = Chroma( 31 | persist_directory=persist_directory, 32 | embedding_function=self.embedding_function, 33 | client_settings=chromadb.config.Settings( 34 | anonymized_telemetry=False, 35 | is_persistent=True 36 | ) 37 | ) 38 | logging.info("[__init__] Loaded existing vector store from disk") 39 | except Exception as e: 40 | logging.error(f"[__init__] Error loading existing vector store: {e}") 41 | self.db = None 42 | else: 43 | self.db = None 44 | else: 45 | self.db = None 46 | logging.info("[__init__] Dry run mode - skipping DB load") 47 | 48 | def initialize_from_pdf(self, pdf_path, chunk_size=300, chunk_overlap=50): 49 | """Initialize vector DB with PDF content.""" 50 | if self.dry_run: 51 | logging.info(f"[initialize_from_pdf] Dry run mode - skipping PDF initialization for {pdf_path}") 52 | return True 53 | 54 | logging.debug(f"[initialize_from_pdf] Loading PDF: {pdf_path}") 55 | try: 56 | loader = PyPDFLoader(pdf_path) 57 | documents = loader.load() 58 | logging.debug(f"[initialize_from_pdf] Loaded {len(documents)} pages from PDF") 59 | 60 | text_splitter = CharacterTextSplitter( 61 | chunk_size=chunk_size, 62 | chunk_overlap=chunk_overlap, 63 | separator="\n" 64 | ) 65 | docs = text_splitter.split_documents(documents) 66 | logging.debug(f"[initialize_from_pdf] Split into {len(docs)} chunks") 67 | 68 | self.db = Chroma.from_documents( 69 | documents=docs, 70 | embedding=self.embedding_function, 71 | persist_directory=self.persist_directory, 72 | client_settings=chromadb.config.Settings( 73 | anonymized_telemetry=False, 74 | is_persistent=True 75 | ) 76 | ) 77 | logging.debug("[initialize_from_pdf] Vector store initialized successfully") 78 | return True 79 | 80 | except Exception as e: 81 | logging.error(f"[initialize_from_pdf] Error initializing vector store: {e}") 82 | return False 83 | 84 | def search(self, query, n_results=3): 85 | """Search the vector store for relevant documents.""" 86 | if self.dry_run: 87 | logging.info(f"[search] Dry run mode - returning mock results for query: {query}") 88 | return "Mock result: This is a dry run response." 89 | 90 | logging.debug(f"[search] Searching for: {query}") 91 | try: 92 | if not self.db: 93 | logging.error("[search] Vector store not initialized") 94 | return "" 95 | 96 | docs = self.db.similarity_search(query, k=n_results) 97 | context = "\n".join(doc.page_content for doc in docs) 98 | logging.debug(f"[search] Found {len(docs)} relevant documents") 99 | return context 100 | 101 | except Exception as e: 102 | logging.error(f"[search] Error searching vector store: {e}") 103 | return "" 104 | 105 | def __del__(self): 106 | """Cleanup when the object is destroyed.""" 107 | logging.debug("[__del__] Vector store cleanup complete") --------------------------------------------------------------------------------