enhanced graph

.gitignore

@@ -3,4 +3,5 @@
 data
 __pycache__/
 logs/
-test.py
+test.py
+reports/

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import os
-from application.services import streamlit_function, gemini_model
+from application.services import gemini_api_service, streamlit_function
 from google.genai.errors import ClientError
 from application.utils import logger
 from application.schemas.response_schema import (
@@ -44,7 +44,7 @@ if st.session_state.pdf_file:
         with col1:
             if st.button(f"Generate {MODEL_1} Response"):
                 with st.spinner(f"Calling {MODEL_1}..."):
-                    result = gemini_model.extract_emissions_data_as_json(API_1 , MODEL_1, st.session_state.pdf_file[0],FULL_RESPONSE_SCHEMA)
+                    result = gemini_api_service.extract_emissions_data_as_json(API_1 , MODEL_1, st.session_state.pdf_file[0],FULL_RESPONSE_SCHEMA)
                     excel_file = streamlit_function.export_results_to_excel(result, MODEL_1, file_name)
                     st.session_state[f"{MODEL_1}_result"] = result
             if st.session_state[f"{MODEL_1}_result"]:
@@ -54,7 +54,7 @@ if st.session_state.pdf_file:
         with col2:
             if st.button(f"Generate {MODEL_2} Response"):
                 with st.spinner(f"Calling {MODEL_2}..."):
-                    result = gemini_model.extract_emissions_data_as_json(API_2, MODEL_2, st.session_state.pdf_file[0],FULL_RESPONSE_SCHEMA)
+                    result = gemini_api_service.extract_emissions_data_as_json(API_2, MODEL_2, st.session_state.pdf_file[0],FULL_RESPONSE_SCHEMA)
                     excel_file = streamlit_function.export_results_to_excel(result, MODEL_2, file_name)
                     st.session_state[f"{MODEL_2}_result"] = result
             if st.session_state[f"{MODEL_2}_result"]:
@@ -65,7 +65,7 @@ if st.session_state.pdf_file:
             try:
                 if st.button(f"Generate {MODEL_3} Response"):
                     with st.spinner(f"Calling {MODEL_3}..."):
-                        result = gemini_model.extract_emissions_data_as_json(API_3, MODEL_3, st.session_state.pdf_file[0], FULL_RESPONSE_SCHEMA)
+                        result = gemini_api_service.extract_emissions_data_as_json(API_3, MODEL_3, st.session_state.pdf_file[0], FULL_RESPONSE_SCHEMA)
                         excel_file = streamlit_function.export_results_to_excel(result, MODEL_3, file_name)
                         st.session_state[f"{MODEL_3}_result"] = result
             except ClientError as e:

application/services/{gemini_model.py → gemini_api_service.py}

@@ -5,6 +5,8 @@ from typing import Optional, Dict, Union, IO, List, BinaryIO
 from google import genai
 from google.genai import types
 from application.utils import logger
+import requests
+import io
 
 logger=logger.get_logger()
 
@@ -136,11 +138,11 @@ def upload_file(
     config: Optional[Dict[str, str]] = None
 ) -> Optional[types.File]:
     """
-    Uploads a file to the Gemini API, handling both file paths and binary streams.
+    Uploads a file to the Gemini API, handling local file paths, binary streams, and URLs.
 
     Args:
-        file (Union[str, IO[bytes]]): File path or binary file object (e.g., from Streamlit).
-        file_name (Optional[str]): Name for the file. If None, attempts to use file.name.
+        file (Union[str, IO[bytes]]): Local file path, URL, or binary file object.
+        file_name (Optional[str]): Name for the file. If None, tries to infer it from the source.
         config (Optional[Dict[str, str]]): Extra config like 'mime_type'.
 
     Returns:
@@ -150,8 +152,14 @@ def upload_file(
         Exception: If upload fails.
     """
     try:
+        # Determine if input is a URL
+        is_url = isinstance(file, str) and file.startswith(('http://', 'https://'))
+
+        # Determine file name if not provided
         if not file_name:
-            if isinstance(file, str):
+            if is_url:
+                file_name = os.path.basename(file.split("?")[0])  # Remove query params
+            elif isinstance(file, str):
                 file_name = os.path.basename(file)
             elif hasattr(file, "name"):
                 file_name = os.path.basename(file.name)
@@ -164,17 +172,32 @@ def upload_file(
         config.update({"name": sanitized_name, "mime_type": mime_type})
         gemini_file_key = f"files/{sanitized_name}"
 
+        # Check if file already exists
         if gemini_file_key in get_files():
             logger.info(f"File already exists on Gemini: {gemini_file_key}")
             return client.files.get(name=gemini_file_key)
 
         logger.info(f"Uploading file to Gemini: {gemini_file_key}")
 
+        # Handle URL
+        if is_url:
+            headers = {
+                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
+            }
+            response = requests.get(file, headers=headers)
+            response.raise_for_status()
+            file_content = io.BytesIO(response.content)
+            return client.files.upload(file=file_content, config=config)
+
+        # Handle local file path
         if isinstance(file, str):
+            if not os.path.isfile(file):
+                raise FileNotFoundError(f"Local file '{file}' does not exist.")
             with open(file, "rb") as f:
                 return client.files.upload(file=f, config=config)
-        else:
-            return client.files.upload(file=file, config=config)
+
+        # Handle already opened binary file object
+        return client.files.upload(file=file, config=config)
 
     except Exception as e:
         logger.error(f"Failed to upload file '{file_name}': {e}")

application/services/mongo_db_service.py

@@ -84,4 +84,5 @@ def retrieve_documents(collection_name: str, query: Optional[Dict] = None) -> Li
         logger.exception(f"An error occurred while retrieving documents: {str(e)}")
         return []
     
-# all_docs = retrieve_documents("Zalando")
+# all_docs = retrieve_documents("Zalando")
+# print(all_docs)

application/tools/emission_data_extractor.py

@@ -6,7 +6,7 @@ import requests
 from google import genai
 from google.genai import types
 from application.utils.logger import get_logger
-from application.services.gemini_model import upload_file
+from application.services.gemini_api_service import upload_file
 from application.services.mongo_db_service import store_document
 from application.schemas.response_schema import GEMINI_GHG_PARAMETERS
 from langchain_core.tools import tool

application/tools/web_search_tools.py

@@ -8,6 +8,7 @@ from typing import Literal
 from duckduckgo_search import DDGS
 from tavily import TavilyClient
 from langchain_core.tools import tool
+import ast
 
 logger = get_logger()
 load_dotenv()
@@ -54,7 +55,8 @@ def get_top_companies_from_web(query: str):
 
         output = response.output_text
         # logger.info(f"Raw Output: {output}")
-        parsed_list = eval(output.strip())
+        parsed_list = ast.literal_eval(output.strip())
+        # parsed_list = eval(output.strip())
         logger.info(f"Parsed List: {parsed_list}")
         result =  CompanyListResponse(companies=parsed_list)
         return result

main.py

@@ -0,0 +1,161 @@
+import os
+import operator
+import functools
+from typing import Annotated, Sequence, TypedDict, Union, Optional
+
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables import Runnable
+from langchain.output_parsers.openai_tools import JsonOutputKeyToolsParser
+from langgraph.graph import StateGraph, END
+
+from application.agents.scraper_agent import scraper_agent
+from application.agents.extractor_agent import extractor_agent
+from application.utils.logger import get_logger
+
+load_dotenv()
+logger = get_logger()
+
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    logger.error("OPENAI_API_KEY is missing. Please set it in your environment variables.")
+    raise EnvironmentError("OPENAI_API_KEY not found in environment variables.")
+
+MEMBERS = ["Scraper", "Extractor"]
+OPTIONS = ["FINISH"] + MEMBERS
+
+SUPERVISOR_SYSTEM_PROMPT = (
+    "You are a supervisor tasked with managing a conversation between the following workers: {members}. "
+    "Given the user's request and the previous messages, determine what to do next:\n"
+    "- If the user asks to search, find, or scrape data from the web, choose 'Scraper'.\n"
+    "- If the user asks to extract ESG emissions data from a file or PDF, choose 'Extractor'.\n"
+    "- If the task is complete, choose 'FINISH'.\n"
+    "- If the message is general conversation (like greetings, questions, thanks, chatting), directly respond with a message.\n"
+    "Each worker will perform its task and report back.\n"
+    "When you respond directly, make sure your message is friendly and helpful."
+)
+
+FUNCTION_DEF = {
+    "name": "route_or_respond",
+    "description": "Select the next role OR respond directly.",
+    "parameters": {
+        "title": "RouteOrRespondSchema",
+        "type": "object",
+        "properties": {
+            "next": {
+                "title": "Next Worker",
+                "anyOf": [{"enum": OPTIONS}],
+                "description": "Choose next worker if needed."
+            },
+            "response": {
+                "title": "Supervisor Response",
+                "type": "string",
+                "description": "Respond directly if no worker action is needed."
+            }
+        },
+        "required": [],
+    },
+}
+
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+    next: Optional[str]
+    response: Optional[str]
+
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+def agent_node(state: AgentState, agent: Runnable, name: str) -> dict:
+    logger.info(f"Agent {name} invoked.")
+    try:
+        result = agent.invoke(state)
+        logger.info(f"Agent {name} completed successfully.")
+        return {"messages": [HumanMessage(content=result["output"], name=name)]}
+    except Exception as e:
+        logger.exception(f"Agent {name} failed with error: {str(e)}")
+        raise
+
+prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", SUPERVISOR_SYSTEM_PROMPT),
+        MessagesPlaceholder(variable_name="messages"),
+        (
+            "system",
+            "Based on the conversation, either select next worker (one of: {options}) or respond directly with a message.",
+        ),
+    ]
+).partial(options=str(OPTIONS), members=", ".join(MEMBERS))
+
+# supervisor_chain = (
+#     prompt
+#     | llm.bind_functions(functions=[FUNCTION_DEF], function_call="route_or_respond")
+#     | JsonOutputFunctionsParser()
+# )
+
+supervisor_chain = (
+    prompt
+    | llm.bind_tools(tools=[FUNCTION_DEF], tool_choice="route_or_respond")
+    | JsonOutputKeyToolsParser(key_name="route_or_respond")
+)
+
+def supervisor_node(state: AgentState) -> AgentState:
+    logger.info("Supervisor invoked.")
+    output = supervisor_chain.invoke(state)
+    logger.info(f"Supervisor output: {output}")
+
+    if isinstance(output, list) and len(output) > 0:
+        output = output[0]  
+
+    next_step = output.get("next")
+    response = output.get("response")
+
+    if not next_step and not response:
+        raise ValueError(f"Supervisor produced invalid output: {output}")
+
+    return {
+        "messages": state["messages"],
+        "next": next_step,
+        "response": response,
+    }
+
+workflow = StateGraph(AgentState)
+
+workflow.add_node("Scraper", functools.partial(agent_node, agent=scraper_agent, name="Scraper"))
+workflow.add_node("Extractor", functools.partial(agent_node, agent=extractor_agent, name="Extractor"))
+workflow.add_node("supervisor", supervisor_node)
+# workflow.add_node("supervisor", supervisor_chain)
+workflow.add_node("supervisor_response", lambda state: {"messages": [AIMessage(content=state["response"], name="Supervisor")]})
+
+for member in MEMBERS:
+    workflow.add_edge(member, "supervisor")
+
+def router(state: AgentState):
+    if state.get("response"):
+        return "supervisor_response"
+    return state.get("next")
+
+conditional_map = {member: member for member in MEMBERS}
+conditional_map["FINISH"] = END
+conditional_map["supervisor_response"] = "supervisor_response"
+
+
+workflow.add_conditional_edges("supervisor", router, conditional_map)
+
+workflow.set_entry_point("supervisor")
+
+graph = workflow.compile()
+
+# # === Example Run ===
+if __name__ == "__main__":
+    logger.info("Starting the graph execution...")
+    initial_message = HumanMessage(content="Can you get zalando pdf link")
+    input_state = {"messages": [initial_message]}
+
+    for step in graph.stream(input_state):
+        if "__end__" not in step:
+            logger.info(f"Graph Step Output: {step}")
+            print(step)
+            print("----")
+
+    logger.info("Graph execution completed.")

pages/chatbot.py

@@ -2,16 +2,10 @@ import streamlit as st
 from dotenv import load_dotenv
 
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
-# from application.agents.scraper_agent import app
-# from application.utils.logger import get_logger
 
-try:
-    from application.agents.scraper_agent import app
-    # from application.main import graph
-    from application.utils.logger import get_logger
-except ImportError as e:
-    st.error(f"Import Error: Ensure backend modules are accessible. Details: {e}")
-    st.stop()
+from application.agents.scraper_agent import app
+from main import graph
+from application.utils.logger import get_logger
 
 logger = get_logger()
 
@@ -19,8 +13,8 @@ st.set_page_config(page_title="Sustainability AI Assistant", layout="wide")
 st.title("♻️ Sustainability Report AI Assistant")
 st.caption(
     "Ask about sustainability reports by company or industry! "
-    "(e.g., 'Get report for Apple', 'Download report for Microsoft 2023', "
-    "'Find reports for top 3 airline companies', 'Download this pdf <link>')"
+    "(e.g., 'Get sustainability report for Apple', 'Download sustainability report for Microsoft 2023', "
+    "'Find sustainability reports for top 3 airline companies', 'Download this pdf <link>')"
 )
 
 load_dotenv()
@@ -34,10 +28,10 @@ def initialize_chat_history():
 def display_chat_history():
     """Render previous chat messages."""
     for message in st.session_state.messages:
-        if isinstance(message, SystemMessage):
-            # st.info(f"System: {message.content}")
-            pass
-        elif isinstance(message, HumanMessage):
+        # if isinstance(message, SystemMessage):
+        #     # st.info(f"System: {message.content}")
+        #     pass
+        if isinstance(message, HumanMessage):
             with st.chat_message("user"):
                 st.markdown(message.content)
         elif isinstance(message, AIMessage):
@@ -77,10 +71,10 @@ def display_last_ai_response():
         logger.warning("No AI message found in the final output.")
 
 initialize_chat_history()
-display_chat_history()
 
 if user_query := st.chat_input("Your question about sustainability reports..."):
     logger.info(f"User input received: {user_query}")
+    display_chat_history()
 
     st.session_state.messages.append(HumanMessage(content=user_query))
 

pages/multiple_pdf_extractor.py

@@ -6,7 +6,7 @@ from application.schemas.response_schema import (
     GEMINI_GOVERNANCE_PARAMETERS, GEMINI_MATERIALITY_PARAMETERS,
     GEMINI_NET_ZERO_INTERVENTION_PARAMETERS
 )
-from application.services import streamlit_function, gemini_model
+from application.services import gemini_api_service, streamlit_function
 from application.utils import logger
 
 logger = logger.get_logger()
@@ -58,7 +58,7 @@ if st.session_state.uploaded_files:
                         all_results = {}
                         
                         for label, schema in RESPONSE_SCHEMAS.items():
-                            result = gemini_model.extract_emissions_data_as_json("gemini", selected_model, pdf_file, schema)
+                            result = gemini_api_service.extract_emissions_data_as_json("gemini", selected_model, pdf_file, schema)
                             streamlit_function.export_results_to_excel(result, sheet_name=selected_model, filename=file_name, column=label)
                             all_results[label] = result
                         st.session_state[result_key] = all_results