Integrated multi-agent workflow from llama index.

.gitignore

@@ -0,0 +1,3 @@
+.env
+/chat
+/code

app.py

@@ -1,5 +1,103 @@
+import os
 import gradio as gr
-from huggingface_hub import InferenceClient
+from llama_index.core.tools import FunctionTool
+from llama_index.core.agent.workflow import AgentWorkflow, ReActAgent
+from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
+from code_agent import initialize_code_agent
+from scientific_paper_agent import load_scientific_paper_dataset, ScientificPaperRetriever
+from search_agent import init_search_tool
+from tools import math_tool_func, init_image_to_text
+from web_agent import initialize_web_agent
+
+global currentMode
+
+hf_token = os.environ.get('HF_TOKEN')
+
+llm = HuggingFaceInferenceAPI(
+    model_name="Qwen/Qwen2.5-7B-Instruct",
+    token=hf_token
+)
+image_to_text_tool = FunctionTool.from_defaults(
+    fn=init_image_to_text,
+    name="image_to_text_tool",
+    description="Generate captions from an image URL using BLIP. Returns both conditional and unconditional captions."
+)
+search_tool = init_search_tool()
+math_tool = FunctionTool.from_defaults(
+    fn=math_tool_func,
+    name="math_tool",
+    description="Solving math problems using the Qwen2.5-Math-1.5B model."
+)
+
+scientific_paper_dataset = load_scientific_paper_dataset()
+scientific_paper_tool = FunctionTool.from_defaults(
+    fn=ScientificPaperRetriever(scientific_paper_dataset).run,
+    name="scientific_paper_info_retriever",
+    description="Retrieves detailed information about scientific papers."
+)
+
+# Define Agents
+code_agent = initialize_code_agent()
+image_to_text_agent = ReActAgent(
+    name="image_to_text",
+    description="Generate text captions from images",
+    tools=[image_to_text_tool],
+    system_prompt=(
+        "You are an assistant specialized in image understanding. "
+        "When given an image URL, use the image_to_text_tool to generate captions. "
+        "Provide both conditional and unconditional descriptions in clear, concise language. "
+        "Do not invent details beyond what the tool provides."
+    ),
+    llm=llm
+)
+math_agent = ReActAgent(
+    name="math_solver",
+    description="Solve math problems using a dedicated math model",
+    tools=[math_tool],
+    system_prompt=(
+        "You are an assistant specialized in solving math problems. "
+        "When given a math query, use the math_solver_tool to compute the answer. "
+        "Explain the solution clearly and step by step when possible, "
+        "but keep the final answer concise and accurate."
+    ),
+    llm=llm
+)
+search_web_agent = ReActAgent(
+    name="search_web",
+    description="Searches the web for answers",
+    tools=[search_tool],
+    system_prompt=(
+        "You are a helpful assistant. Use DuckDuckGoSearch to look up information. "
+        "Always summarize the first useful result and return it directly. "
+        "Do not keep searching repeatedly."
+    ),
+    llm=llm
+)
+scientific_paper_agent = ReActAgent(
+    name="scientific_paper_agent",
+    description="Search scientific papers for the agent",
+    tools=[scientific_paper_tool],
+    system_prompt="You are a helpful assistant that can answer scientific questions based on scientific papers.",
+    llm=llm
+)
+query_engine_agent = initialize_web_agent(llm)
+
+
+# DEFINE THE WORKFLOW
+multi_agent_workflow = AgentWorkflow(
+    agents=[
+        query_engine_agent,
+        search_web_agent,
+        math_agent,
+        image_to_text_agent,
+        scientific_paper_agent,
+        code_agent
+    ],
+    root_agent="query_engine",
+    initial_state={ "num_of_calls": 0 },
+    state_prompt="Current state: {state}. User Message: {msg}"
+)
+
 
 
 def respond(
@@ -9,62 +107,44 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-    hf_token: gr.OAuthToken,
+    mode
 ):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
-
-    messages = [{"role": "system", "content": system_message}]
-
-    messages.extend(history)
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+    global currentMode
+    print("Current Mode: " + mode)
+
+    if mode == "Math Mode":
+        currentMode = "math"
+    elif mode == "Conversation Mode":
+        currentMode = "conversation"
+    elif mode == "Image Mode":
+        currentMode = "image"
+    else:
+        currentMode = "conversation"
+
+    yield multi_agent_workflow.run(message, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
+
 
 with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+    # Dropdown placed above the chat input
+    mode_dropdown = gr.Dropdown(
+        choices=["Math Mode", "Conversation Mode", "Image Mode"],
+        value="Conversation Mode",
+        label="Interaction Mode"
+    )
+
+    # ChatInterface without additional_inputs
+    chatbot = gr.ChatInterface(
+        fn=respond,
+        type="messages"
+    )
 
+    # Link dropdown value to respond function
+    mode_dropdown.change(
+        lambda m: m,
+        inputs=mode_dropdown,
+        outputs=[]
+    )
 
 if __name__ == "__main__":
     demo.launch()
+

code_agent.py

@@ -0,0 +1,49 @@
+import os
+import chromadb
+from llama_index.core import VectorStoreIndex
+from llama_index.core.tools import QueryEngineTool
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
+# from llama_index.llms.litellm import LiteLLM
+from llama_index.core.agent.workflow import ReActAgent
+
+def initialize_code_agent():
+    hf_token = os.environ.get('HF_TOKEN')
+    deepseek_token = os.environ.get('DEEPSEEK_TOKEN')
+
+    code_db = chromadb.PersistentClient(path="./code_db")
+
+    code_chroma_collection = code_db.get_or_create_collection('code')
+    code_vector_store = ChromaVectorStore(chroma_collection=code_chroma_collection)
+
+    embedding_model = HuggingFaceEmbedding(
+        model_name="BAAI/bge-small-en-v1.5",
+        device="cpu",
+        token=hf_token,
+    )
+    index = VectorStoreIndex.from_vector_store(code_vector_store, embed_model=embedding_model)
+
+    code_llm = HuggingFaceInferenceAPI(
+        model_name="deepseek-ai/deepseek-coder-1.3b-instruct",
+        api_key=deepseek_token,
+        token=hf_token,
+    )
+    code_query_engine = index.as_query_engine(
+        llm=code_llm,
+        similarity_top_k=3
+    )
+    code_query_engine_tool = QueryEngineTool.from_defaults(
+        query_engine=code_query_engine,
+        name="my_code_query_engine",
+        description="Code Query engine for the agent",
+        return_direct=False
+    )
+    return ReActAgent(
+        name="code_engine",
+        description="Query engine for the agent",
+        tools=[code_query_engine_tool],
+        system_prompt="You are a calculator assistant. Use your tools for any math operation.",
+        llm=code_llm
+    )
+

requirements.txt

@@ -0,0 +1,20 @@
+accelerate
+datasets
+smolagents
+llama-index
+huggingface_hub
+llama-index-llms-huggingface-api
+langchain_core
+langchain_community
+llama-index-embeddings-huggingface
+llama-index-tools-duckduckgo
+rank_bm25
+chromadb
+llama-index-vector-stores-chroma
+torch
+torchvision
+torchaudio
+pillow
+transformers
+llama-index-llms-litellm
+llama-index-utils-workflow

scientific_paper_agent.py

@@ -0,0 +1,46 @@
+import datasets
+from langchain_core.documents import Document
+from langchain_community.retrievers import BM25Retriever
+
+def load_scientific_paper_dataset():
+    # Convert dataset entries into Document objects
+    scientific_paper_dataset = datasets.load_dataset("gsasikiran/Summarize-Scientific-Papers-Processed", split="train")
+    docs = [
+        Document(
+            page_content="\n".join([
+                f"Title: {scientific_paper['title']}",
+                f"Authors: {scientific_paper['authors']}",
+                f"What is it: {scientific_paper['article_classification']}",
+                f"Claims: {scientific_paper['claims']}",
+                f"Contradictions: {scientific_paper['contradictions_and_limitations']}",
+                f"Ethical Considerations: {scientific_paper['ethical_considerations']}",
+                f"Summary: {scientific_paper['executive_summary']}",
+                f"Subfield: {scientific_paper['field_subfield']}",
+                f"Theorical Implications: {scientific_paper['interpretation_and_theoretical_implications']}",
+                f"Method to Retrieve Info: {scientific_paper['methodological_details']}",
+                f"People used to get data: {scientific_paper['procedures_and_architectures']}",
+                f"Context of Research: {scientific_paper['research_context']}",
+                f"Research Hypothesis: {scientific_paper['research_question_and_hypothesis']}",
+                f"Three Takeways: {scientific_paper['three_takeaways']}",
+                f"Type of Paper: {scientific_paper['type_of_paper']}"
+            ]),
+            metadata={"title": scientific_paper["title"]}
+        )
+        for scientific_paper in scientific_paper_dataset
+    ]
+    return docs
+
+
+# --- Scientific Paper Retriever ---
+class ScientificPaperRetriever:
+    def __init__(self, docs):
+        # Build BM25 retriever from documents
+        self.retriever = BM25Retriever.from_documents(docs)
+
+    def run(self, query: str) -> str:
+        results = self.retriever.retrieve(query)
+        if results:
+            return "\n\n".join([doc.text for doc in results[:3]])
+        else:
+            return "No matching scientific paper found."
+

search_agent.py

@@ -0,0 +1,7 @@
+from llama_index.core.tools import FunctionTool
+from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
+
+def init_search_tool():
+    search_tool_spec = DuckDuckGoSearchToolSpec()
+    search_tool = FunctionTool.from_defaults(search_tool_spec.duckduckgo_full_search)
+    return search_tool

tools.py

@@ -0,0 +1,59 @@
+import requests
+import os
+import torch
+from PIL import Image
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from llama_index.core.tools import FunctionTool
+
+hf_token = os.environ.get("HF_TOKEN")
+
+# Load processor and model once (outside the function for efficiency)
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
+
+math_model_id = "Qwen/Qwen2.5-Math-1.5B"
+math_tokenizer = AutoTokenizer.from_pretrained(math_model_id, use_auth_token=hf_token)
+math_model = AutoModelForCausalLM.from_pretrained(
+    math_model_id,
+    dtype=torch.float16,
+    device_map="auto",
+    use_auth_token=hf_token
+)
+
+
+def math_tool_func(problem: str) -> str:
+    inputs = math_tokenizer(problem, return_tensors="pt").to(math_model.device)
+    outputs = math_model.generate(**inputs, max_new_tokens=128)
+    result = math_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return result
+
+
+
+def init_image_to_text(img_url: str) -> dict:
+    """
+    Convert an image URL into text captions using BLIP.
+
+    Args:
+        img_url (str): URL of the image to caption.
+
+    Returns:
+        dict: Contains both conditional and unconditional captions.
+    """
+    raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+
+    # Conditional captioning
+    conditional_prompt = "a photography of"
+    inputs_cond = processor(raw_image, conditional_prompt, return_tensors="pt")
+    out_cond = model.generate(**inputs_cond)
+    conditional_caption = processor.decode(out_cond[0], skip_special_tokens=True)
+
+    # Unconditional captioning
+    inputs_uncond = processor(raw_image, return_tensors="pt")
+    out_uncond = model.generate(**inputs_uncond)
+    unconditional_caption = processor.decode(out_uncond[0], skip_special_tokens=True)
+
+    return {
+        "conditional_caption": conditional_caption,
+        "unconditional_caption": unconditional_caption,
+    }

web_agent.py

@@ -0,0 +1,40 @@
+import os
+import chromadb
+from llama_index.core import VectorStoreIndex
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from llama_index.core.tools import QueryEngineTool
+
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
+from llama_index.core.agent.workflow import ReActAgent
+
+def initialize_web_agent(llm: HuggingFaceInferenceAPI):
+    hf_token = os.environ.get('HF_TOKEN')
+    db = chromadb.PersistentClient(path="./chat_db")
+    chroma_collection = db.get_or_create_collection("chat")
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+
+    embedding_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5", device="cpu")
+    index = VectorStoreIndex.from_vector_store(vector_store, embed_model=embedding_model)
+
+    query_engine = index.as_query_engine(
+        llm=llm,
+        similarity_top_k=3
+    )
+
+    query_engine_tool = QueryEngineTool.from_defaults(
+        query_engine=query_engine,
+        name="my_query_engine",
+        description="Query engine for the agent",
+        return_direct=False
+    )
+
+    return ReActAgent(
+        name="query_engine",
+        description="Query engine for the agent",
+        tools=[query_engine_tool],
+        system_prompt="You are a calculator assistant. Use your tools for any math operation.",
+        llm=llm
+    )
+
+