Upload 6 files

src/app.py

@@ -0,0 +1,130 @@
+import streamlit as st
+from run_model import generate_response, generate_RAG_response
+from web_search import search_web
+from deep_research import perform_deep_research
+import tempfile
+
+st.set_page_config(layout="wide")
+
+def main():
+    st.title("💬 Chat with Gemma")
+    
+
+    with st.sidebar:
+        st.title("Tools")
+        st.markdown("For source code of this website, visit : <some_link>")
+        
+        option = st.selectbox(
+            "Choose tools",
+            ("Simple Chat", "Web Search", "Upload PDF", "Deep Web Search"),
+        )
+
+        temperature = st.slider(
+            label="Temperature (controls randomness)",
+            min_value=0.0,
+            max_value=2.0,
+            value=1.0,
+            step=0.01,
+            help="Lower = more deterministic, Higher = more random"
+        )
+        # Top-k sampling
+        top_k = st.slider(
+            label="Top-k (limits to top K tokens by probability)",
+            min_value=0,
+            max_value=100,
+            value=50,
+            step=1,
+            help="0 = disable top-k filtering"
+        )
+
+        # Top-p (nucleus sampling)
+        top_p = st.slider(
+            label="Top-p (nucleus sampling cutoff)",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.9,
+            step=0.01,
+            help="0.0 = conservative, 1.0 = more random"
+        )
+
+        if option == "Upload PDF":
+            file = st.file_uploader(label="Uploaded file will provide context to LLM", type="pdf")
+            
+
+        if option == "Web Search":
+            st.write("Web Search Enabled for next query")
+            # WB_SEARCH = True
+
+        if option == "None":
+            st.warning("You are not using any tool")
+
+        if option == "Deep Web Search":
+            st.write("Deep Web Research Enabled for next query")
+
+
+
+    # col1, col2 = st.columns([6, 1], gap="small")
+
+    # column 1
+    # with col1:
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    for msg in st.session_state.messages:
+        with st.chat_message(msg["role"]):
+            st.markdown(msg["content"])
+
+    # prompt = st.chat_input("Say something...")
+
+    if prompt:= st.chat_input("Say something..."):
+        # Display user message
+        st.chat_message("user").markdown(prompt)
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        if option == "Simple Chat":
+
+            response = generate_response(history=st.session_state.messages, query=prompt, temperature=temperature, top_k=top_k, top_p=top_p)
+            st.chat_message("assistant").markdown(response)
+            # st.session_state.messages.append({"role": "assistant", "content": response})
+        
+        if option == "Web Search":
+            st.session_state.messages.append([{"role": "user", "content" : prompt}])
+
+            response, sources = search_web(prompt)
+            # asnswer = response
+            # st.chat_message("assistant").markdown(f"{response}\n\n###Sources\n{'\n'.join([source for source in sources])}")
+            with st.chat_message("assistant"):
+                st.markdown(f"{response}\n\n### Sources\n" + "\n".join(sources))
+
+        if option == "Deep Web Search":
+            st.session_state.messages.append([{"role": "user", "content" : prompt}])
+
+            response = perform_deep_research(prompt)
+            # asnswer = response
+            # st.chat_message("assistant").markdown(f"{response}\n\n###Sources\n{'\n'.join([source for source in sources])}")
+            with st.chat_message("assistant"):
+                st.markdown(response)
+
+        if option == "Upload PDF":
+            st.session_state.messages.append([{"role": "user", "content" : prompt}])
+
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
+                tmp_file.write(file.read())
+                tmp_path = tmp_file.name
+            print(tmp_path)
+            response = generate_RAG_response(prompt, tmp_path, st.session_state.messages)
+            # print(file)
+            with st.chat_message("assistant"):
+                st.markdown(response)
+        
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        
+        
+    # # Slider in column 2
+    # with col2:
+    #     temperature = st.slider("Temperature", 0.0, 1.0, 0.0)
+    #     top_k = st.slider("Top k", 0.0, 100.0, 40.0)
+    #     top_p = st.slider("Top p", 0.0, 1.0, 0.95)
+
+
+if __name__ == "__main__":
+    main()

src/deep_research.py

@@ -0,0 +1,296 @@
+from langchain_google_genai import ChatGoogleGenerativeAI
+import os
+from typing import Annotated, Dict, Any, List, Union, Optional
+from typing_extensions import TypedDict, Annotated, Literal
+import operator
+from dataclasses import dataclass, field
+from re import T
+import json
+from langchain_core.runnables import RunnableConfig
+from langgraph.graph import START, END, StateGraph
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+import time
+from tavily import TavilyClient
+import re
+from model import get_gemma
+
+from api_key import TAVILY_API_KEY
+
+gemma_model = get_gemma()
+
+tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
+
+max_web_research_loops: int = 4
+
+
+@dataclass(kw_only=True)
+class SummaryState:
+    research_topic: str = field(default=None)
+    search_query: str = field(default=None)
+    web_search_results: Annotated[list, operator.add] = field(default_factory=list)
+    sources_gathered: Annotated[list, operator.add] = field(default=list)
+    research_loop_count: int = 0
+    running_summary: str = None
+
+
+@dataclass(kw_only=True)
+class SummaryStateInput:
+    research_topic: str = None
+
+@dataclass(kw_only=True)
+class SummaryStateOutput:
+    running_summary: str = None
+
+
+# Query Writer
+query_writer_instructions = """Your gola is to generate web search query.
+The query will gather information about specific topic.
+
+Topic: {research_topic}
+
+Return your query as JSON object:
+{{
+    "query": "string",
+    "aspect" : "string",
+    "rationale" : "string"
+}}
+"""
+
+# Summerizer Instructions
+summerizer_instructions = """Your goal is to generate high-quality summary of the web search results.
+when EXTENDING an existing summary:
+1. Seamlessly integrate new information without repeating what's already covered.
+2. Maintain consistancy with existing content's style.
+3. Only add new and non-redudant information.
+4. Ensure smooth transition between existing and new content.
+
+when creating a NEW summary:
+1. Highlight the most relevant information from each source.
+2. Provide concise overview of the key points related to each report topic.
+3. Emphasize on significant findings or insights.
+4. Ensure coherent flow of information.
+
+In both cases:
+1. Focus on factual & objective information
+2. Maintain consistat technical depth
+3. Avoid repetition & redundancy
+4. DON'T use phrases like "based on new results"
+5. DON'T add preamble like "Here is an extended summary ...", instead just provide summary directly
+6. DON'T add references or works cited section.
+7. You will generate tables using markdown when user asks you to do.
+"""
+
+# Reflection Instructions
+reflection_summary = """You are an expert research assistant analyzing summary about {research_topic}.
+
+Your Tasks :
+1. Identify knowledge gaps or areas the need further exploration.
+2. Generate a follow-up question that would help in expanding the understanding.
+3. Focus on technical details, implementation specifics.
+
+Ensure follow-up question is self-contained and includes necessary context for web search.
+
+Return response as JSON object:
+{{
+    "knowledge_gap" : "string",
+    "follow_up_query" : "string"
+}}
+"""
+
+
+def generate_query(state: SummaryState):
+    # To generate query for web search
+
+    system_message_for_query_writer = query_writer_instructions.format(research_topic=state.research_topic)
+
+    result = gemma_model.invoke(
+        [
+            HumanMessage(content=f"IMPORTANT INSTRUCTIONS:\n{system_message_for_query_writer}\n\nGenerate a query for web search")
+        ]
+    )
+
+    # print(f"[FUN] GENERATE_QUERY:\nType: {type(result)}\nContent: {result}")
+    raw_content = result.content.strip()
+    match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", raw_content, re.DOTALL)
+    if match:
+        json_str = match.group(1)
+    else:
+        raise ValueError(f"Failed to extract JSON from model response: {raw_content}")
+
+    query = json.loads(json_str)
+    return {"search_query" : query["query"]}
+
+
+def deduplicate_and_format_sources(
+    search_response: Union[Dict[str, Any], List[Dict[str, Any]]],
+    max_tokens_per_source: int,
+    fetch_full_page: bool = False
+) -> str:
+    """
+    Format and deduplicate search responses from various search APIs.
+
+    Takes either a single search response or list of responses from search APIs,
+    deduplicates them by URL, and formats them into a structured string.
+
+    Args:
+        search_response (Union[Dict[str, Any], List[Dict[str, Any]]]): Either:
+            - A dict with a 'results' key containing a list of search results
+            - A list of dicts, each containing search results
+        max_tokens_per_source (int): Maximum number of tokens to include for each source's content
+        fetch_full_page (bool, optional): Whether to include the full page content. Defaults to False.
+
+    Returns:
+        str: Formatted string with deduplicated sources
+
+    Raises:
+        ValueError: If input is neither a dict with 'results' key nor a list of search results
+    """
+    # Convert input to list of results
+    if isinstance(search_response, dict):
+        sources_list = search_response['results']
+    elif isinstance(search_response, list):
+        sources_list = []
+        for response in search_response:
+            if isinstance(response, dict) and 'results' in response:
+                sources_list.extend(response['results'])
+            else:
+                sources_list.extend(response)
+    else:
+        raise ValueError("Input must be either a dict with 'results' or a list of search results")
+
+    # Deduplicate by URL
+    unique_sources = {}
+    for source in sources_list:
+        if source['url'] not in unique_sources:
+            unique_sources[source['url']] = source
+
+    # Format output
+    formatted_text = "Sources:\n\n"
+    for i, source in enumerate(unique_sources.values(), 1):
+        formatted_text += f"Source: {source['title']}\n===\n"
+        formatted_text += f"URL: {source['url']}\n===\n"
+        formatted_text += f"Most relevant content from source: {source['content']}\n===\n"
+        if fetch_full_page:
+            # Using rough estimate of 4 characters per token
+            char_limit = max_tokens_per_source * 4
+            # Handle None raw_content
+            raw_content = source.get('raw_content', '')
+            if raw_content is None:
+                raw_content = ''
+                print(f"Warning: No raw_content found for source {source['url']}")
+            if len(raw_content) > char_limit:
+                raw_content = raw_content[:char_limit] + "... [truncated]"
+            formatted_text += f"Full source content limited to {max_tokens_per_source} tokens: {raw_content}\n\n"
+
+    return formatted_text.strip()
+
+
+def format_sources(search_results: Dict[str, Any]) -> str:
+    """
+    Format search results into a bullet-point list of sources with URLs.
+
+    Creates a simple bulleted list of search results with title and URL for each source.
+
+    Args:
+        search_results (Dict[str, Any]): Search response containing a 'results' key with
+                                        a list of search result objects
+
+    Returns:
+        str: Formatted string with sources as bullet points in the format "* title : url"
+    """
+    return '\n'.join(
+        f"* {source['title']} : {source['url']}"
+        for source in search_results['results']
+    )
+
+
+def web_research(state: SummaryState):
+    search_results = tavily_client.search(state.search_query, include_raw_content=True, max_results=1)
+    search_str = deduplicate_and_format_sources(search_results, max_tokens_per_source=1000)
+    return {
+        "sources_gathered" : [format_sources(search_results)],
+        "research_loop_count" : state.research_loop_count + 1,
+        "web_search_results" : [search_str]
+    }
+
+
+def summarize_sources(state: SummaryState):
+    existing_summary = state.running_summary
+    print(state.web_search_results)
+    most_recent_web_search = state.web_search_results[-1]
+
+    if existing_summary:
+        human_message = (
+            f"IMPORTANT INSTRUCTIONS:\n{summerizer_instructions}\n\n"
+            f"Extend the existing summary: {existing_summary}\n\n"
+            f"Include new search results: {most_recent_web_search}"
+            f"That addresses the following topic: {state.research_topic}"
+        )
+    else:
+        human_message = (
+            f"IMPORTANT INSTRUCTIONS:\n{summerizer_instructions}\n\n"
+            f"Generate summary of these search results: {most_recent_web_search}"
+            f"That addresses the following topic: {state.research_topic}"
+        )
+
+    result = gemma_model.invoke([HumanMessage(content=human_message)])
+
+    return {"running_summary" : result.content}
+
+
+def reflect_on_summary(state: SummaryState):
+    result = gemma_model([
+        HumanMessage(content=f"IMPORTANT INSTRUCTIONS:\n{reflection_summary.format(research_topic=state.research_topic)}\n\nIdentify a knowledge gap and generate a follow-up web search query based on existing knowledge: {state.running_summary}")
+    ])
+#    print(f">> [FUN] REFLECT ON SUMMARY:\n{result.content}")
+    raw_content = result.content.strip()
+    match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", raw_content, re.DOTALL)
+    if match:
+        json_str = match.group(1)
+    else:
+        raise ValueError(f"Failed to extract JSON from model response: {raw_content}")
+
+    query = json.loads(json_str)
+    print(query)
+    return {"search_query" : query["follow_up_query"]}
+
+
+def finalize_summary(state: SummaryState):
+    all_sources = "\n".join(source for source in state.sources_gathered)
+    print(f"All Sources: {all_sources}")
+    running_summary = f"## Summary\n\n{state.running_summary}\n\nSources:\n{all_sources}"
+    return {"running_summary" : running_summary}
+
+def route_research(state: SummaryState):
+    if state.research_loop_count <= max_web_research_loops:
+        return "web_research"
+    else:
+        return "finalize_summary"
+    
+
+def perform_deep_research(query):
+    
+    builder = StateGraph(SummaryState, input_schema=SummaryStateInput, output_schema=SummaryStateOutput)
+
+    builder.add_node("generate_query", generate_query)
+    builder.add_node("web_research", web_research)
+    builder.add_node("summarize_sources", summarize_sources)
+    builder.add_node("reflect_on_summary", reflect_on_summary)
+    builder.add_node("finalize_summary", finalize_summary)
+
+    # Add edges
+    builder.add_edge(START, "generate_query")
+    builder.add_edge("generate_query", "web_research")
+    builder.add_edge("web_research", "summarize_sources")
+    builder.add_edge("summarize_sources", "reflect_on_summary")
+    builder.add_conditional_edges("reflect_on_summary", route_research)
+    builder.add_edge("finalize_summary", END)
+
+
+    graph = builder.compile()
+
+    research_input = SummaryStateInput(research_topic=query)
+
+    research_output = graph.invoke(research_input)
+
+    return research_output["running_summary"]

src/model.py

@@ -0,0 +1,19 @@
+from api_key import GOOGLE_API_KEY
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+
+
+def get_gemma(temperature=0, top_k=40, top_p=0.95):
+
+    gemma_model = ChatGoogleGenerativeAI(
+            model="gemma-3-12b-it",
+            temperature=temperature,
+            max_tokens=2048,
+            timeout=None,
+            max_retries=2,
+            google_api_key=GOOGLE_API_KEY,
+            top_k=top_k,
+            top_p=top_p 
+    )
+
+    return gemma_model

src/rag_utils.py

@@ -0,0 +1,54 @@
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+
+
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+def load_pdf_document(file_path):
+    document_loader = PyPDFLoader(file_path)
+    return document_loader.load()
+
+def chunk_documents(raw_documents):
+    text_processor = RecursiveCharacterTextSplitter(
+        chunk_size = 1000,
+        chunk_overlap = 200,
+        add_start_index = True
+    )
+    return text_processor.split_documents(raw_documents)
+
+def find_related_documents(query, vector_database):
+    # return vector_database.similarity_search(query, k=2)
+    return vector_database.max_marginal_relevance_search(query, k=2, fetch_k=5, lambda_mult=0.6)
+
+
+def ProcessDocuments(document_path: str) -> str:
+
+    loaded_doc = load_pdf_document(document_path)
+    chunked_doc = chunk_documents(loaded_doc)
+
+    
+    vector_database = Chroma(
+        persist_directory=f"./chroma_store/{document_path.split("\\")[-1].split(".")[0]}",
+        embedding_function=embedding_model
+    )
+
+    vector_database.add_documents(chunked_doc)
+
+
+def generate_context(query: str, file: str):
+
+    ProcessDocuments(file)
+
+    
+    vector_database = Chroma(
+        persist_directory=f"./chroma_store/{file.split("\\")[-1].split(".")[0]}",
+        embedding_function=embedding_model
+    )
+
+    relevant_docs = find_related_documents(query, vector_database)
+    context_text = "\n".join([doc.page_content for doc in relevant_docs])
+ 
+    return query, context_text
+

src/run_model.py

@@ -0,0 +1,21 @@
+from model import get_gemma
+from rag_utils import generate_context
+
+
+def generate_response(history=[], temperature: float=0.0, top_k=None, top_p=None):
+    
+    gemma_model = get_gemma(temperature=temperature, top_k=top_k, top_p=top_p)
+   
+    response = gemma_model.invoke(history).content
+
+    return response
+
+def generate_RAG_response(query: str, file_path, history=[]):
+    gemma_model = get_gemma()
+    query, context = generate_context(query, file_path)
+
+    history[-1] = {"role" : "user", "content" : f"INSTRUCTION: Answer the query with given context in mind.\nQUERY: {query}\n\nCONTEXT : {context}"}
+    
+    response = gemma_model.invoke(history).content
+
+    return response

src/web_search.py

@@ -0,0 +1,76 @@
+from langchain.agents import AgentExecutor, create_react_agent
+from langchain.agents import load_tools, Tool
+# from langchain.tools import DuckDuckGoSearchResults
+from tavily import TavilyClient
+from langchain.prompts import PromptTemplate
+from model import get_gemma
+from api_key import TAVILY_API_KEY
+
+gemma_model = get_gemma()
+
+
+# Create the ReAct template
+react_template = """Answer the following questions as best you can. You have access to the following tools:
+
+{tools}
+
+Use the following format:
+
+Question: the input question you must answer
+Thought: you should always think about what to do
+Action: the action to take, should be one of [{tool_names}]
+Action Input: the input to the action
+Observation: the result of the action
+... (this Thought/Action/Action Input/Observation can repeat N times)
+Thought: I now know the final answer
+Final Answer: the final answer to the original input question
+
+Begin!
+
+Question: {input}
+Thought:{agent_scratchpad}"""
+
+prompt = PromptTemplate(
+    template=react_template,
+    input_variables=["tools", "tool_names", "input", "agent_scratchpad"]
+)
+
+
+tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
+
+tavily_search_tool = Tool(
+    name="tavily search",
+    description = "A web search engine. Use this to as a search engine for general queries.",
+    func = lambda x: tavily_client.search(x, max_results=1)
+)
+
+# Prepare tools
+tools = load_tools(["llm-math"], llm=gemma_model)
+tools.append(tavily_search_tool)
+
+
+# Construct the ReAct agent
+agent = create_react_agent(gemma_model, tools, prompt)
+agent_executor = AgentExecutor(
+    agent=agent,
+    tools=tools,
+    verbose=True,
+    handle_parsing_errors=True,
+    return_intermediate_steps=True
+)
+
+def get_urls_from_response(response):
+    urls = []
+    for step in response["intermediate_steps"]:
+        urls.append(step[1]["results"][0]["url"])
+    return urls
+
+
+def search_web(query):
+    response = agent_executor.invoke({"input" : query})
+
+    output = response["output"]
+    sources = get_urls_from_response(response)
+
+    return output, sources 
+