Upload 16 files

agent.py

@@ -1,224 +1,77 @@
-"""LangGraph Agent"""
-import os
-from dotenv import load_dotenv
-from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import tools_condition
-from langgraph.prebuilt import ToolNode
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_groq import ChatGroq
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.document_loaders import ArxivLoader
-from langchain_community.vectorstores import SupabaseVectorStore
-from langchain_core.messages import SystemMessage, HumanMessage
-from langchain_core.tools import tool
-from langchain.tools.retriever import create_retriever_tool
-from supabase.client import Client, create_client
-from langchain_community.tools import DuckDuckGoSearchRun
-
-load_dotenv()
-
-@tool
-def search_answer(question: str) -> str:
-    """
-    This function uses the DuckDuckGoSearchRun tool to perform a search.
-    """
-    search = DuckDuckGoSearchRun()
+# agent.py
+from langchain_ollama.chat_models import ChatOllama
+from langchain_core.messages import HumanMessage, ToolMessage
+from tools.wiki import wikipedia_search_tool
+import re, ast
+from configs.config import Config
+from configs.registry import TOOL_REGISTRY
+
+# Define Variables
+env = Config()
+llm = env.LOCAL_LLM
+
+tools_registery = TOOL_REGISTRY
+
+def generate_prompt(query: str) -> str:
+    tool_list = "\n".join(
+        f"- {name}: {meta['description']}" for name, meta in tools_registery.items()
+    )
 
-    return search.invoke(question)
+    return f"""
+    You are a smart assistant that decides which tool to use based on user queries.
 
-@tool
-def multiply(a: int, b: int) -> int:
-    """Multiply two numbers.
-    Args:
-        a: first int
-        b: second int
-    """
-    return a * b
-
-@tool
-def add(a: int, b: int) -> int:
-    """Add two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
-    return a + b
-
-@tool
-def subtract(a: int, b: int) -> int:
-    """Subtract two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
-    return a - b
-
-@tool
-def divide(a: int, b: int) -> int:
-    """Divide two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
-    if b == 0:
-        raise ValueError("Cannot divide by zero.")
-    return a / b
-
-@tool
-def modulus(a: int, b: int) -> int:
-    """Get the modulus of two numbers.
-    
-    Args:
-        a: first int
-        b: second int
+    User Query: "{query}"
+
+    Available tools:
+    {tool_list}
+
+    Respond in this format:
+    Tool: [tool_name]
+    Tool Input: [Python dict of parameters]
     """
-    return a % b
-
-@tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
-    
-    Args:
-        query: The search query."""
-    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"wiki_results": formatted_search_docs}
-
-@tool
-def web_search(query: str) -> str:
-    """Search Tavily for a query and return maximum 3 results.
-    
-    Args:
-        query: The search query."""
-    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"web_results": formatted_search_docs}
-
-@tool
-def arvix_search(query: str) -> str:
-    """Search Arxiv for a query and return maximum 3 result.
-    
-    Args:
-        query: The search query."""
-    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"arvix_results": formatted_search_docs}
-
-
-
-# load the system prompt from the file
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    system_prompt = f.read()
-
-# System message
-sys_msg = SystemMessage(content=system_prompt)
-
-# build a retriever
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") #  dim=768
-supabase: Client = create_client(
-    os.environ.get("SUPABASE_URL"), 
-    os.environ.get("SUPABASE_SERVICE_KEY"))
-vector_store = SupabaseVectorStore(
-    client=supabase,
-    embedding= embeddings,
-    table_name="documents",
-    query_name="match_documents_langchain",
-)
-create_retriever_tool = create_retriever_tool(
-    retriever=vector_store.as_retriever(),
-    name="Question Search",
-    description="A tool to retrieve similar questions from a vector store.",
-)
-
-
-
-tools = [
-    multiply,
-    add,
-    subtract,
-    divide,
-    modulus,
-    wiki_search,
-    web_search,
-    arvix_search,
-    search_answer,
-]
-
-# Build graph function
-def build_graph(provider: str = "groq"):
-    """Build the graph"""
-    # Load environment variables from .env file
-    if provider == "google":
-        # Google Gemini
-        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    elif provider == "groq":
-        # Groq https://console.groq.com/docs/models
-        llm = ChatGroq(model="qwen-qwq-32b", temperature=0) # optional : qwen-qwq-32b gemma2-9b-it
-    elif provider == "huggingface":
-        # TODO: Add huggingface endpoint
-        llm = ChatHuggingFace(
-            llm=HuggingFaceEndpoint(
-                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
-                temperature=0,
-            ),
-        )
-    else:
-        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
-    # Bind tools to LLM
-    llm_with_tools = llm.bind_tools(tools)
-
-    # Node
-    def assistant(state: MessagesState):
-        """Assistant node"""
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
-    def retriever(state: MessagesState):
-        """Retriever node"""
-        similar_question = vector_store.similarity_search(state["messages"][0].content)
-        example_msg = HumanMessage(
-            content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
-        )
-        return {"messages": [sys_msg] + state["messages"] + [example_msg]}
-
-    builder = StateGraph(MessagesState)
-    builder.add_node("retriever", retriever)
-    builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
-    builder.add_edge(START, "retriever")
-    builder.add_edge("retriever", "assistant")
-    builder.add_conditional_edges(
-        "assistant",
-        tools_condition,
-    )
-    builder.add_edge("tools", "assistant")
 
-    # Compile graph
-    return builder.compile()
+def parse_tool_selection(response: str) -> tuple[str, dict]:
+    tool_match = re.search(r"Tool:\s*(\w+)", response)
+    input_match = re.search(r"Tool Input:\s*(\{.*\})", response)
+
+    if not tool_match or not input_match:
+        raise ValueError("Failed to parse tool selection.")
+
+    tool_name = tool_match.group(1)
+    tool_input = ast.literal_eval(input_match.group(1))
+    return tool_name, tool_input
+
+def main(query: str = None):
+    user_query = query.strip()
+
+    # 1. Generate selection prompt
+    prompt = generate_prompt(user_query)
+    response = llm.invoke([HumanMessage(content=prompt)])
+
+    # 2. Parse tool selection
+    try:
+        tool_name, tool_input = parse_tool_selection(response.content)
+    except Exception as e:
+        print("Error parsing tool selection:", e)
+        print("LLM response was:", response.content)
+        return
+
+    # 3. Run selected tool
+    tool_entry = tools_registery.get(tool_name)
+    if not tool_entry:
+        print(f"Tool '{tool_name}' not found.")
+        return
+
+    tool = tool_entry["tool"]
+
+    try:
+        result = tool.invoke(tool_input)
+    except Exception as e:
+        print(f"Error running tool '{tool_name}': {e}")
+        return
+
+    print("Final Answer:", result.content)
 
-# test
 if __name__ == "__main__":
-    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
-    # Build the graph
-    graph = build_graph(provider="huggingface")
-    # Run the graph
-    messages = [HumanMessage(content=question)]
-    messages = graph.invoke({"messages": messages})
-    for m in messages["messages"]:
-        m.pretty_print()
+    query = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of English Wikipedia."
+    main(query)

configs/config.py

@@ -0,0 +1,76 @@
+import os
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+from llama_index.embeddings.openai import OpenAIEmbedding
+from langchain_ollama import ChatOllama, OllamaEmbeddings
+
+# Load environment variables from .env file
+load_dotenv()
+
+class Config:
+    LLM_PROVIDER="ollama"
+    if LLM_PROVIDER == "ollama":
+        # Ollama configuration
+        OLLAMA_BASE_URL="http://localhost:11434"
+        #LOCAL_LLM_MODEL="deepseek-r1:8b"
+        #LOCAL_LLM_MODEL="deepseek-r1:7b"
+        LOCAL_LLM_MODEL = "llama3.2"
+        LOCAL_LLM = ChatOllama(model=LOCAL_LLM_MODEL, 
+                               base_url=OLLAMA_BASE_URL,
+                               temperature=0.5)
+        EMBED_MODEL = OllamaEmbeddings(model="nomic-embed-text")
+
+    elif LLM_PROVIDER == "openai":
+        OPENAI_API_KEY:str = os.getenv("OPENAI_API_KEY","")
+        LLM_MODEL_NAME:str = os.getenv("LLM_MODEL","gpt-3.5-turbo")
+        LLM_VIDEO_MODEL_NAME:str = os.getenv("LLM_VIDEO_MODEL","gpt-4o-mini")
+
+        LLM = ChatOpenAI(model=LLM_MODEL_NAME, openai_api_key=OPENAI_API_KEY)
+        EMBED_MODEL = OpenAIEmbedding(openai_api_key=OPENAI_API_KEY)
+
+    FILE = None
+
+    WIKI_DEFAULT_PROMPTS = {
+        "system": (
+            "You are an intelligent assistant with access to Wikipedia search results related to the user's query.\n"
+            "Use only the information provided in the search results to answer the question accurately.\n"
+            "Carefully analyze the query to determine what the user is asking.\n"
+            "Respond clearly and concisely, avoiding speculation or information not found in the provided content.\n"
+            "If the answer is not present in the search results, state that explicitly."
+        ),
+        "user": "{query}"
+        }
+
+    SYSTEM_MSG = f"""
+    You are a general-purpose AI assistant.
+
+    When I ask you a question:
+    - Think step by step to determine the answer.
+    - List your reasoning steps clearly.
+    - If additional information is required to answer the question, use the 'wiki' tool by providing the directive: tool_call: [wiki].
+    - Provide your final output using one of the following formats:
+    - FINAL ANSWER: [YOUR FINAL ANSWER]
+    - tool_call: [TOOL_NAME]
+
+    Only one of these should appear in your final output — either FINAL ANSWER or tool_call.
+
+    If you are unsure or need more information, always use the 'wiki' tool.
+
+    Final answer formatting rules:
+    - If the answer is a number:
+        - Do NOT use commas (e.g., write 1000 not 1,000).
+        - Do NOT include units like "$" or "%" unless explicitly requested.
+    - If the answer is a string:
+        - Do NOT use articles (e.g., "a", "an", "the").
+        - Do NOT use abbreviations (e.g., write "New York" instead of "NY").
+        - Write digits as plain text (e.g., "four" instead of "4") unless stated otherwise.
+    - If the answer is a comma-separated list:
+        - Follow the same rules above for each item depending on whether it’s a number or a string.
+
+    Do NOT include any additional arguments in tool calls.
+
+    Available tools:
+    - wikipedia_search_tool: Search Wikipedia.
+    """
+
+config = Config()

configs/registry.py

@@ -0,0 +1,9 @@
+
+from tools.wiki import wikipedia_search_tool
+
+TOOL_REGISTRY = {
+    "wikipedia_search_tool": {
+        "tool": wikipedia_search_tool,
+        "description": "Searches Wikipedia for information. Use when user asks about a topic, person, or event.",
+    },
+}

tools/wiki.py

@@ -0,0 +1,112 @@
+
+import sys
+import os
+from langchain_core.tools import tool
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+from utils.call_llm import llm
+from configs.config import Config
+env = Config()
+
+def generate_search_string(query: str) -> str:
+    """
+    Generate an optimal Wikipedia search string from the given query.
+
+    Args:
+        query (str): The input query for generating the search string.
+
+    Returns:
+        str: A single continuous search string optimized for Wikipedia search.
+    """
+    if not query or not isinstance(query, str):
+        raise ValueError("Query must be a non-empty string.")
+
+    prompt = f"""
+    Generate an optimal Wikipedia search string from the query '{query}'. \n
+    Just return a single continuous search string without any additional text or formatting or quotation marks. \n
+    Do not include any other text or explanation."""
+
+    response = env.LOCAL_LLM.invoke(prompt)
+    if not response or not response.content.strip():
+        raise ValueError("Failed to generate a valid search string.")
+    
+    return response.content.strip()
+
+def document_store(query, chunk_size, chunk_overlap):
+    """Load a Wikipedia page based on the query and language."""
+    from langchain_community.document_loaders import WikipediaLoader
+    from langchain.text_splitter import RecursiveCharacterTextSplitter
+    from langchain.schema.document import Document
+    from langchain_community.vectorstores.faiss import FAISS
+
+    embedding_model = env.EMBED_MODEL
+    language = "en"
+
+    search_query = generate_search_string(query)
+    if not search_query:
+        raise ValueError("Search query is empty or invalid.")
+        
+    loader = WikipediaLoader(query=search_query, lang=language)
+    documents = loader.load()
+    combined_text = "".join([doc.page_content for doc in documents if doc.page_content])
+    if not combined_text:
+        raise ValueError("No text found in the loaded documents.")
+
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+    )
+        
+    chunks = splitter.split_text(combined_text)
+    if not chunks:
+        raise ValueError("No chunks generated from the combined text.")
+
+    docs = [
+        Document(page_content=chunk, metadata={"source": query})
+        for chunk in chunks
+    ]
+    if not docs:
+        raise ValueError("No documents created from the chunks.")
+
+    embeddings = embedding_model.embed_documents([doc.page_content for doc in docs])
+    if not embeddings:
+        raise ValueError("No embeddings generated for the documents.")
+
+    store = FAISS.from_documents(docs, embedding=embedding_model)
+    return store
+
+def search(query,chunk_size, chunk_overlap):
+    store = document_store(query,chunk_size, chunk_overlap)
+    results = store.similarity_search_with_score(query, k=5)
+
+    # Filter results based on a relevance threshold
+    filtered_results = []
+    for doc, score in results:
+        if score <= 0.5:  # Relevance threshold
+            filtered_results.append((doc, score))
+    
+    return filtered_results
+    
+@tool("wikipedia_search_tool")
+def wikipedia_search_tool(query: str, chunk_size: int =1000, chunk_overlap: int =200):
+    """
+    Run the Wikipedia search tool with the given query and parameters.
+    """
+    print("----- Wiki Run ---")
+    default_prompts = env.WIKI_DEFAULT_PROMPTS
+
+    response = search(query, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    if not response:
+        response = [("No relevant documents found.", 1.0)]
+
+    llm_input = [
+        {"role": "system", "content": default_prompts["system"]},
+        {"role": "user", "content": default_prompts["user"].format(query=query)},
+        {"role": "user", "content": response[0][0] if response else "No relevant documents found."}
+    ]
+
+    call_llm = env.LOCAL_LLM.invoke(llm_input)
+
+    return call_llm
+
+

tools/youtube.py

@@ -0,0 +1,103 @@
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+class YoutubeSearchTool:
+    def __init__(self, query: str, chunk_size: int = 1000, chunk_overlap: int = 200):
+        from configs.config import Config
+        from utils.generate_search_string import generate_search_string
+        
+        env = Config()
+        self.generate_search_string = generate_search_string
+        self.llm = env.LLM_VIDEO_MODEL_NAME
+        self.embedding_model = env.EMBED_MODEL
+        self.query = query
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.language = "en"
+
+    def extract_youtube_link(self) -> str:
+        """ Generate a YouTube search URL based on the query """
+        import re
+
+        youtube_url_pattern = r"https?://www\.youtube\.com/watch\?v=[\w-]+"
+        match = re.search(youtube_url_pattern, self.query)
+    
+        return match.group(0) if match else None
+
+    def video_loader(self):
+        """ Load a YouTube video based on the query and language """
+
+        from langchain_community.document_loaders import YoutubeLoader
+        from langchain_community.document_loaders.youtube import TranscriptFormat
+        from langchain.text_splitter import RecursiveCharacterTextSplitter
+        from langchain.schema import Document
+
+        url = self.extract_youtube_link()
+        optimized_string = self.generate_search_string(self.query)
+
+        loader = YoutubeLoader.from_youtube_url(
+            url,
+            add_video_info=True,
+            transcript_format=TranscriptFormat.CHUNKS,
+            chunk_size_seconds=30,
+            language=self.language,
+            )
+
+        documents = loader.load()
+
+        # combined_text = "".join([doc.page_content for doc in documents])
+
+        # # Split into chunks
+        # splitter = RecursiveCharacterTextSplitter(
+        #     chunk_size=self.chunk_size,
+        #     chunk_overlap=self.chunk_overlap,
+        # )
+        # chunks = splitter.split_text(combined_text)
+        
+        # return chunks
+    
+    def vector_store(self):
+        """ Create a vector store from the video chunks """
+        from langchain_community.vectorstores import FAISS
+        from langchain_openai import OpenAIEmbeddings
+
+        chunks = self.video_loader()
+        if not chunks:
+            return "No relevant video chunks found."
+
+        docs = [
+            Document(page_content=chunk, metadata={"source": self.query})
+            for chunk in chunks
+        ]
+        self.vector_store = FAISS.from_documents(docs, embedding = OpenAIEmbeddings())
+        
+    def run(self, query: str):
+        """ Run the YouTube search tool with the given query """
+        print("----- YouTube Search Tool Run ---")
+
+        store = self.vector_store()
+
+        results = store.similarity_search_with_score(query, k=1)
+                        
+        if not results:
+            return "No relevant video chunks found in the vector store."
+        
+        return [
+            {
+                "content": doc.page_content,
+                "score": score,
+                "source": doc.metadata.get("source", "Unknown")
+            }
+            for doc, score in results
+        ]
+
+
+if __name__ == "__main__":
+    # Example usage
+    query = "Examine the video at https://www.youtube.com/watch?v=1htKBjuUWec. What does Teal'c say in response to the question 'Isn't that hot?'"
+    
+    youtube_tool = YoutubeSearchTool(query=query)
+    youtube_tool.run(query)
+    print(f"Search URL: {youtube_tool.extract_youtube_link(query)}")
+    print("Video chunks loaded successfully.")

utils/call_llm.py

@@ -0,0 +1,50 @@
+from langchain_ollama import ChatOllama
+from langchain_core.messages import SystemMessage
+
+def llm(
+    model_name: str,
+    messages: list[dict],
+    temperature: float = 0.1,
+    max_tokens: int = 1024,
+    **kwargs,
+) -> str:
+    """
+    Calls the Ollama Chat model and returns the generated response content.
+    """
+    try:
+        print("[NODE] ----- Calling Ollama Chat -----")
+
+        # Construct the prompt with explicit separation for SystemMessage
+        prompt_parts = []
+        for message in messages:
+            if isinstance(message, dict):
+                prompt_parts.append(f"{message['role'].capitalize()}: {message['content']}")
+            elif isinstance(message, SystemMessage):
+                prompt_parts.append(f"System: {message.content}")
+            else:
+                prompt_parts.append(message)
+
+        prompt = "\n\n".join(prompt_parts)  # Add extra separation for clarity
+
+        print(f"Constructed Prompt:\n{prompt}")
+
+        chat = ChatOllama(
+            model=model_name,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            **kwargs
+        )
+
+        response = chat.invoke(prompt)
+        print("----- Ollama Chat response -----")
+        print(response.content)
+
+        if not response or not response.content:
+            print("No content returned from the Ollama Chat model.")
+            return "No content generated."
+
+        return response.content
+
+    except Exception as e:
+        print("An error occurred while calling the Ollama Chat model: %s", str(e))
+        raise