Upload 10 files

agent.py

@@ -0,0 +1,264 @@
+import json
+from typing import Any, Literal
+
+from langgraph.graph import END, StateGraph
+from langchain_community.document_loaders import JSONLoader
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+
+from langgraph.store.memory import InMemoryStore
+from langchain.chat_models.base import BaseChatModel
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_community.utilities import SQLDatabase
+
+from constant import DB_PATH
+from state import AgentState
+import logging
+
+# Initialize logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class SQLAgentRAG:
+    def __init__(
+        self,
+        llm: BaseChatModel,
+        tools: Any,
+        db_uri: str = DB_PATH,
+        table_json_path: str = "table.jsonl",
+        column_json_path: str = "column.jsonl"
+    ):
+        self.llm = llm
+        self.table_json_path = table_json_path
+        self.column_json_path = column_json_path
+        self.db = SQLDatabase.from_uri(db_uri)
+        self.schema = None
+        self.retriever = tools
+
+         # add nodes
+        graph = StateGraph(AgentState)
+        graph.add_node("router", self.router_node)
+        graph.add_node("general_asistant", self._general_asistant)
+        graph.add_node("sql_gen", self._sql_gen)
+        graph.add_node("validate_sql", self._validate_sql)
+        graph.add_node("solve_error", self._solve_error)
+        graph.add_node("response", self._query_gen_node)
+
+        # add edges
+        graph.set_entry_point("router")
+        graph.add_edge("sql_gen", "validate_sql")
+        graph.add_conditional_edges(
+            "router",
+            self.router,
+            {
+                "SQL": "sql_gen",
+                "GENERAL": "general_asistant"
+            }
+        )
+        graph.add_conditional_edges(
+            "validate_sql",
+            self._should_continue
+        )
+        graph.add_edge("solve_error", "validate_sql")
+        graph.add_edge("response", END)
+        graph.add_edge("general_asistant", END)
+
+        # compile
+        store = InMemoryStore()
+        checkpointer = MemorySaver()
+        self.graph = graph.compile(checkpointer=checkpointer, store=store)
+
+
+    def _indexing_table(self, query: str):
+        """
+        Index and retrieve relevant tables based on the input query.
+        """
+        logger.info("Indexing Table...")
+        docs_table = JSONLoader(
+            file_path=self.table_json_path,
+            jq_schema='.',
+            text_content=False,
+            json_lines=True
+        ).load()
+
+        retriever = self.retriever(docs_table, k=5, search_type='mmr', lambda_mult=1)
+
+        matched_documents_table = retriever.invoke(query)
+        matched_tables = [
+            json.loads(doc.page_content)["table"] for doc in matched_documents_table
+        ]
+
+        return matched_tables
+    
+    def _indexing_column(self, matched_tables, query: str):
+        """
+        Index and retrieve relevant columns based on the matched tables and query.
+        """
+        logger.info("Get matched schema...")
+        docs_column = JSONLoader(
+            file_path=self.column_json_path,
+            jq_schema='.',
+            text_content=False,
+            json_lines=True
+        ).load()
+
+        retriever = self.retriever(docs_column, k=20, search_type='similarity')
+
+        matched_columns = retriever.invoke(query)
+        matched_columns_filtered = [
+            json.loads(doc.page_content) for doc in matched_columns
+            if json.loads(doc.page_content)["table_name"] in matched_tables
+        ]
+
+        matched_columns_cleaned = [
+            f'table_name={doc["table_name"]}|column_name={doc["column_name"]}|data_type={doc["data_type"]}'
+            for doc in matched_columns_filtered
+        ]
+
+        return matched_columns_cleaned
+    
+    def _sql_gen(self, state: AgentState):
+        """
+        Generates a SQL query based on the input provided by the user.
+        This function uses the LLM to construct the query from matched tables and columns.
+        """
+        logger.info("Generate SQL Query...")
+        messages = state["messages"][-1].content
+        matched_table = self._indexing_table(messages)
+        self.schema = self._indexing_column(matched_table, messages)
+
+        prompt = PromptTemplate(
+            template="""
+              You are a SQL master expert specializing in writing complex SQL queries for SQLite. Your task is to construct a SQL query based on the provided information. Follow these strict rules:
+
+              QUERY: {query}
+              -------
+              MATCHED_SCHEMA: {matched_schema}
+              -------
+
+              Please construct a SQL query using the MATCHED_SCHEMA and the QUERY provided above.
+              IMPORTANT: Use ONLY the column names (column_name) mentioned in MATCHED_SCHEMA. DO NOT USE any other column names outside of this.
+              IMPORTANT: Associate column_name mentioned in MATCHED_SCHEMA only to the table_name specified under MATCHED_SCHEMA.
+              NOTE: Use SQL 'AS' statement to assign a new name temporarily to a table column or even a table wherever needed.
+              Generate ONLY the SQL query. Do not provide any explanations, comments, or additional text.
+
+            """,
+            input_variables=["query", "matched_schema"]
+        )
+
+        sql_gen = prompt | self.llm | StrOutputParser()
+        result_sql = sql_gen.invoke({"query": messages, "matched_schema": self.schema})
+        return {"sql_query": [result_sql]}
+    
+    def _validate_sql(self, state: AgentState):
+        """
+        Validates the generated SQL query by attempting to execute it.
+        Returns "" if no errors, otherwise returns the error message.
+        """
+        logger.info("Validate SQL...")
+        query = state["sql_query"][-1]
+        try:
+            logger.info(f"Query:\n{query}")
+            self.db.run(query)
+            return {"error_str": ""}
+
+        except Exception as e:
+            return {"error_str": [f"Unexpected Error: {str(e)}"]}
+        
+    def _solve_error(self, state: AgentState):
+        """
+        Called with the error code and error description as the argument to get guidance on how to solve the error
+        """
+        error_string = state.get('error_str', "")
+        sql_query = state.get('sql_query', [None])[-1]
+
+        logger.info(f"Error{error_string}")
+        logger.info(f"SQL query {sql_query}")
+
+        prompt = PromptTemplate(
+            template="""
+              First, identify the main issues with the given SQL query based on the error message.
+              {error_string}
+
+              Next, examine the schema and current SQL query to locate potential sources of the error.
+              {schema}
+
+              Then, modify the current SQL query to fix the error and avoid similar issues in the future.
+              {sql_query}
+
+              Finally, ensure the revised SQL query conforms to the requirements outlined in the original task and provide the corrected SQL query.
+              Generate ONLY the SQL query. Do not provide any explanations, comments, or additional text.
+
+            """,
+            input_variables=["error_string", "schema", "sql_query"]
+        )
+        resolver = prompt | self.llm | StrOutputParser()
+        corrected_query = resolver.invoke({
+            "error_string": error_string, 
+            "schema": self.schema, 
+            "sql_query": sql_query
+        })
+        return {"sql_query": [corrected_query]}
+    
+    def _query_gen_node(self, state: AgentState):
+        """
+        Generates a final response after executing the SQL query and getting the result.
+        """
+        logger.info("Generate Response...")
+        query = state.get('sql_query', [None])[-1]
+        messages = state["messages"][-1]  
+        prompt = PromptTemplate(
+            template="""Based on the following SQL result, generate a natural language response:
+            Query SQL: {user_query}
+            SQL Result: {sql_response}
+            """,
+            input_variables=["user_query", "sql_response"]
+        )
+
+        sql_response = self.db.run(query)
+        gen_llm = prompt | self.llm | StrOutputParser()
+        response = gen_llm.invoke({"user_query": messages, "sql_response": sql_response})
+        return {"messages": [response]}
+    
+    def _general_asistant(self, state: AgentState):
+        """Assistan handle question out of context"""
+        logger.info("Assistant...")
+        messages = state["messages"]
+        response = self.llm.invoke(messages)   
+
+        return {"messages": [response], "sql_query": [""]}
+
+    def router(self, state: AgentState):
+        """Router"""
+        logger.info("Router...")
+        return state["question_type"]
+
+    def router_node(self, state: AgentState):
+        """Router Node"""
+        question = state["messages"][-1].content
+        prompt = PromptTemplate(
+            template= """
+            You are a senior specialist of analytical support. Classify incoming questions into one of two types:
+              - SQL: Related to flight information, schedules, hotels, rentals, recommendations, and anything about vacations
+              - GENERAL: General questions
+            Return only one word: SQL, or GENERAL.
+
+            {question}
+        """,
+            input_variables=["question"]
+        )
+        router = prompt | self.llm | StrOutputParser()
+        question_type = router.invoke({"question": question})
+        return {"question_type": question_type}
+
+    def _should_continue(self, state: AgentState) -> Literal["response", "solve_error"]:
+        """
+        Decides whether to proceed based on SQL validation results.
+        If the last message contains "error_str", 
+        go to error-solving or retry. Otherwise,continue to response. 
+        """
+        if state.get("error_str"):
+            # If error exists, go to solve_error
+            return "solve_error"
+        return "response"

app.py

@@ -0,0 +1,123 @@
+import logging
+import streamlit as st
+from dotenv import load_dotenv, find_dotenv
+from langgraph.errors import GraphRecursionError
+from langchain_groq import ChatGroq
+from agent import SQLAgentRAG
+from tools import retriever
+from constant import GROQ_API_KEY, CONFIG
+
+# Initialize logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+load_dotenv(find_dotenv())
+
+# Initialize the language model
+llm = ChatGroq(
+    model="llama3-8b-8192",
+    api_key=GROQ_API_KEY,
+    temperature=0.1,
+    verbose=True
+)
+
+# Initialize SQL Agent
+agent = SQLAgentRAG(llm=llm, tools=retriever)
+
+def query_rag_agent(query: str):
+    """
+    Handle a query through the RAG Agent, producing an SQL response if applicable.
+
+    Parameters:
+        - query (str): The input query to process.
+
+    Returns:
+        - Tuple[str, List[str]]: The response content and SQL query if applicable.
+
+    Raises:
+        - GraphRecursionError: If there's a recursion limit reached within the agent's graph.
+    """
+    try:
+        output = agent.graph.invoke({"messages": query}, CONFIG)
+        response = output["messages"][-1].content
+        sql_query = output.get("sql_query", ["No SQL query generated"])[-1]
+
+        logger.info(f"Query processed successfully: {query}")
+        return response, sql_query
+
+    except GraphRecursionError:
+        logger.error("Graph recursion limit reached; query processing failed.")
+        return "Graph recursion limit reached. No SQL result generated.", ""
+
+with st.sidebar:
+    st.header("About Project")
+    st.markdown(
+        """
+        RAG (Retrieval-Augmented Generation) Agent SQL is an approach that combines retrieval techniques with text generation to create more relevant and contextualised answers from data, 
+        particularly in SQL databases. RAG-Agent SQL uses two main components:
+        - Retrieval: Retrieving relevant information from the database based on a given question or input.
+        - Augmented Generation: Using natural language models (e.g., LLMs such as GPT or LLaMA) to generate more detailed answers, using information from the retrieval results.
+
+        to see the architecture can be seen here [Github](https://github.com/fahmiaziz98/sql_agent/tree/main/002sql-agent-ra)
+        """
+    )
+
+    st.header("Example Question")
+    st.markdown(
+        """
+        - How many different aircraft models are there? And what are the models?
+        - What is the aircraft model with the longest range?
+        - Which airports are located in the city of Basel?
+        - Can you please provide information on what I asked before?
+        - What are the fare conditions available on Boeing 777-300?
+        - What is the total amount of bookings made in April 2024?
+        - What is the scheduled arrival time of flight number QR0051?
+        - Which car rental services are available in Basel?
+        - Which seat was assigned to the boarding pass with ticket number 0060005435212351?
+        - Which trip recommendations are related to history in Basel?
+        - How many tickets were sold for Business class on flight 30625?
+        - Which hotels are located in Zurich?
+        """
+    )
+
+# Main Application Title
+st.title("RAG SQL-Agent")
+
+# Initialize session state for storing chat messages
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display conversation history from session state
+for message in st.session_state.messages:
+    role = message.get("role", "assistant")
+    with st.chat_message(role):
+        if "output" in message:
+            st.markdown(message["output"])
+        if "sql_query" in message and message["sql_query"]:
+            with st.expander("SQL Query", expanded=True):
+                st.code(message["sql_query"])
+
+# Input form for user prompt
+if prompt := st.chat_input("What do you want to know?"):
+    st.chat_message("user").markdown(prompt)
+    st.session_state.messages.append({"role": "user", "output": prompt})
+
+    # Fetch response from RAG agent function directly
+    with st.spinner("Searching for an answer..."):
+        output_text, sql_query = query_rag_agent(prompt)
+
+    # Display assistant response and SQL query
+    st.chat_message("assistant").markdown(output_text)
+    if sql_query:
+        with st.expander("SQL Query", expanded=True):
+            st.code(sql_query)
+
+    # Append assistant response to session state
+    st.session_state.messages.append(
+        {
+            "role": "assistant",
+            "output": output_text,
+            "sql_query": sql_query,
+        }
+    )

column.jsonl

@@ -0,0 +1,55 @@
+{"table_name":"aircrafts_data","column_name":"aircraft_code","description":"Unique identifier for each aircraft","data_type":"TEXT"}
+{"table_name":"aircrafts_data","column_name":"model","description":"Model of the aircraft","data_type":"TEXT"}
+{"table_name":"aircrafts_data","column_name":"range","description":"Maximum flight range of the aircraft in kilometers","data_type":"INTEGER"}
+{"table_name":"airports_data","column_name":"airport_code","description":"Unique identifier for each airport","data_type":"TEXT"}
+{"table_name":"airports_data","column_name":"airport_name","description":"Full name of the airport","data_type":"TEXT"}
+{"table_name":"airports_data","column_name":"city","description":"City where the airport is located","data_type":"TEXT"}
+{"table_name":"airports_data","column_name":"coordinates","description":"Geographic coordinates of the airport","data_type":"TEXT"}
+{"table_name":"airports_data","column_name":"timezone","description":"Time zone of the airport location","data_type":"TEXT"}
+{"table_name":"boarding_passes","column_name":"ticket_no","description":"Ticket number associated with the boarding pass","data_type":"TEXT"}
+{"table_name":"boarding_passes","column_name":"flight_id","description":"Identifier of the flight","data_type":"INTEGER"}
+{"table_name":"boarding_passes","column_name":"boarding_no","description":"Boarding number assigned to the passenger","data_type":"INTEGER"}
+{"table_name":"boarding_passes","column_name":"seat_no","description":"Seat number assigned to the passenger","data_type":"TEXT"}
+{"table_name":"bookings","column_name":"book_ref","description":"Unique booking reference","data_type":"TEXT"}
+{"table_name":"bookings","column_name":"book_date","description":"Date and time when the booking was made","data_type":"TIMESTAMP"}
+{"table_name":"bookings","column_name":"total_amount","description":"Total amount paid for the booking","data_type":"INTEGER"}
+{"table_name":"car_rentals","column_name":"id","description":"Unique identifier for each car rental","data_type":"INTEGER"}
+{"table_name":"car_rentals","column_name":"name","description":"Name of the car rental company","data_type":"TEXT"}
+{"table_name":"car_rentals","column_name":"location","description":"Location of the car rental","data_type":"TEXT"}
+{"table_name":"car_rentals","column_name":"price_tier","description":"Price category of the rental car","data_type":"TEXT"}
+{"table_name":"car_rentals","column_name":"start_date","description":"Start date of the car rental period","data_type":"DATE"}
+{"table_name":"car_rentals","column_name":"end_date","description":"End date of the car rental period","data_type":"DATE"}
+{"table_name":"car_rentals","column_name":"booked","description":"Booking status of the car (0 for available, 1 for booked)","data_type":"INTEGER"}
+{"table_name":"flights","column_name":"flight_id","description":"Unique identifier for each flight","data_type":"INTEGER"}
+{"table_name":"flights","column_name":"flight_no","description":"Flight number","data_type":"TEXT"}
+{"table_name":"flights","column_name":"scheduled_departure","description":"Scheduled departure date and time","data_type":"TIMESTAMP"}
+{"table_name":"flights","column_name":"scheduled_arrival","description":"Scheduled arrival date and time","data_type":"TIMESTAMP"}
+{"table_name":"flights","column_name":"departure_airport","description":"Code of the departure airport","data_type":"TEXT"}
+{"table_name":"flights","column_name":"arrival_airport","description":"Code of the arrival airport","data_type":"TEXT"}
+{"table_name":"flights","column_name":"status","description":"Current status of the flight","data_type":"TEXT"}
+{"table_name":"flights","column_name":"aircraft_code","description":"Code of the aircraft operating the flight","data_type":"TEXT"}
+{"table_name":"flights","column_name":"actual_departure","description":"Actual departure date and time","data_type":"TIMESTAMP"}
+{"table_name":"flights","column_name":"actual_arrival","description":"Actual arrival date and time","data_type":"TIMESTAMP"}
+{"table_name":"hotels","column_name":"id","description":"Unique identifier for each hotel booking","data_type":"INTEGER"}
+{"table_name":"hotels","column_name":"name","description":"Name of the hotel","data_type":"TEXT"}
+{"table_name":"hotels","column_name":"location","description":"Location of the hotel","data_type":"TEXT"}
+{"table_name":"hotels","column_name":"price_tier","description":"Price category of the hotel","data_type":"TEXT"}
+{"table_name":"hotels","column_name":"checkin_date","description":"Check-in date for the hotel booking","data_type":"DATE"}
+{"table_name":"hotels","column_name":"checkout_date","description":"Check-out date for the hotel booking","data_type":"DATE"}
+{"table_name":"hotels","column_name":"booked","description":"Booking status of the hotel room (0 for available, 1 for booked)","data_type":"INTEGER"}
+{"table_name":"seats","column_name":"aircraft_code","description":"Code of the aircraft","data_type":"TEXT"}
+{"table_name":"seats","column_name":"seat_no","description":"Seat number","data_type":"TEXT"}
+{"table_name":"seats","column_name":"fare_conditions","description":"Fare class of the seat (e.g., Economy, Business)","data_type":"TEXT"}
+{"table_name":"ticket_flights","column_name":"ticket_no","description":"Ticket number","data_type":"TEXT"}
+{"table_name":"ticket_flights","column_name":"flight_id","description":"Flight identifier","data_type":"INTEGER"}
+{"table_name":"ticket_flights","column_name":"fare_conditions","description":"Fare conditions for the ticket (e.g., Economy, Business)","data_type":"TEXT"}
+{"table_name":"ticket_flights","column_name":"amount","description":"Cost of the ticket","data_type":"INTEGER"}
+{"table_name":"tickets","column_name":"ticket_no","description":"Unique ticket number","data_type":"TEXT"}
+{"table_name":"tickets","column_name":"book_ref","description":"Booking reference associated with the ticket","data_type":"TEXT"}
+{"table_name":"tickets","column_name":"passenger_id","description":"Unique identifier for the passenger","data_type":"TEXT"}
+{"table_name":"trip_recommendations","column_name":"id","description":"Unique identifier for each trip recommendation","data_type":"INTEGER"}
+{"table_name":"trip_recommendations","column_name":"name","description":"Name of the recommended attraction or activity","data_type":"TEXT"}
+{"table_name":"trip_recommendations","column_name":"location","description":"Location of the recommended attraction or activity","data_type":"TEXT"}
+{"table_name":"trip_recommendations","column_name":"keywords","description":"Keywords associated with the recommendation","data_type":"TEXT"}
+{"table_name":"trip_recommendations","column_name":"details","description":"Detailed description of the recommendation","data_type":"TEXT"}
+{"table_name":"trip_recommendations","column_name":"booked","description":"Booking status of the recommendation (0 for available, 1 for booked)","data_type":"INTEGER"}

constant.py

@@ -0,0 +1,9 @@
+import os
+GROQ_API_KEY = "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"
+DB_PATH = "sqlite:///travel.sqlite"
+
+CONFIG = {
+    "configurable" : {
+        "thread_id": "1234"
+    }
+}

models.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+from typing import List
+
+class QueryInput(BaseModel):
+    query: str
+
+class QueryOutput(BaseModel):
+    response: str
+    sql_query: List[str]

requirements.txt

@@ -0,0 +1,14 @@
+# Main dependencies
+opentelemetry-api==1.22.0
+pydantic>=2.7
+langchain==0.3.3
+langgraph==0.2.38
+langchain-community==0.3.2
+langchain-groq==0.2.0
+fastembed==0.4.1
+onnx==1.17.0
+onnxruntime==1.18.0
+jq==1.8.0a2
+faiss-cpu==1.8.0
+requests==2.31.0
+streamlit==1.34.0

state.py

@@ -0,0 +1,9 @@
+from typing_extensions import TypedDict
+from langgraph.graph.message import AnyMessage, add_messages
+from typing import Annotated, List
+
+class AgentState(TypedDict):
+    messages: Annotated[List[AnyMessage], add_messages]
+    sql_query: str
+    error_str: str
+    question_type: str

table.jsonl

@@ -0,0 +1,11 @@
+{"table": "aircrafts_data", "description": "This table stores information about aircraft, including aircraft code, model, and flight range.", "example_questions": ["What is the flight range of the aircraft with code '773'?", "Which aircraft model has the shortest flight range?", "What is the difference in flight range between the aircraft with the longest range and the one with the shortest range?"]}
+{"table": "airports_data", "description": "This table contains airport data, including airport code, name, city, coordinates, and time zone.", "example_questions": ["In which city is the airport with code 'ATL' located?", "How many airports are in the Asia time zone?", "Which airport has the northernmost location based on its coordinates?"]}
+{"table": "boarding_passes", "description": "This table stores information about boarding passes, including ticket number, flight ID, boarding number, and seat number.", "example_questions": ["What is the seat number for the ticket number '0060005435212351'?", "How many passengers have seat numbers in row 2 (e.g., 2A, 2B, 2C, etc.)?", "For the flight with flight_id 30625, is there a correlation between boarding number and seat position?"]}
+{"table": "bookings", "description": "This table contains booking information, including booking reference, booking date, and total amount.", "example_questions": ["What is the total amount for the booking with reference '00000F'?", "How many bookings were made in March 2024?", "What is the average booking amount for bookings made in April 2024?"]}
+{"table": "car_rentals", "description": "This table stores information about car rentals, including rental company, location, price tier, rental dates, and booking status.", "example_questions": ["How many car rentals are available in Basel?", "Which rental company offers a luxury car in Basel, and for how many days?", "Calculate the average rental duration for all non-booked cars, grouped by price tier."]}
+{"table": "flights", "description": "This table contains flight information, including flight numbers, schedules, airports, status, and aircraft details.", "example_questions": [ "What is the flight number for the flight departing from BSL to BKK?", "How many flights are scheduled to depart from Shanghai (SHA)?", "For flights that have both scheduled and actual departure times, what is the average delay in minutes?"]}
+{"table": "hotels","description": "This table stores information about hotel bookings, including hotel name, location, price tier, check-in and check-out dates, and booking status.", "example_questions": ["How many hotels are available for booking in Basel?", "Which hotel in Zurich offers an Upscale price tier, and for how many nights?", "Calculate the average length of stay for all non-booked hotels, grouped by price tier."]}
+{"table": "seats", "description": "This table contains information about aircraft seats, including the aircraft code, seat number, and fare conditions.", "example_questions": [ "How many Business class seats are there on aircraft with code '319'?", "List all unique fare conditions available across all aircraft.", "For each aircraft code, what percentage of seats are in the Business fare condition?"]}
+{"table": "ticket_flights", "description": "This table links tickets to specific flights and includes fare conditions and ticket amounts.", "example_questions": ["What is the average ticket amount for Business class on flight_id 30625?", "How many different fare conditions are there across all ticket_flights?", "For each flight_id, what is the price difference between the highest and lowest ticket amount?"]}
+{"table": "tickets", "description": "This table contains information about tickets, including ticket number, booking reference, and passenger ID.", "example_questions": ["How many tickets are associated with the booking reference '06B046'?", "What is the passenger ID for the ticket number '9880005432000987'?", "How many unique passengers (based on passenger_id) have tickets in the system?"]}
+{"table": "trip_recommendations", "description": "This table stores trip recommendations, including name, location, keywords, details, and booking status.", "example_questions": [ "How many trip recommendations are available for Basel?", "What are the keywords associated with the 'Kunstmuseum Basel' recommendation?", "List all unique locations that have trip recommendations, along with the count of recommendations for each location."]}

tools.py

@@ -0,0 +1,38 @@
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
+
+embeddings = FastEmbedEmbeddings()
+
+def retriever(docs, k=5, search_type='mmr', lambda_mult=None):
+    """
+    Creates a document retriever using FAISS (Facebook AI Similarity Search).
+
+    Parameters:
+    -----------
+    docs : List[Document]
+        A list of documents to be indexed for similarity search.
+
+    k : int, optional, default=5
+        The number of top-k results to return for a query.
+
+    search_type : str, optional, default='mmr'
+        The type of search to perform. Options include 'mmr' and 'similarity'.
+
+    lambda_mult : float, optional, default=None
+        Lambda multiplier for Maximal Marginal Relevance (MMR) search.
+
+    Returns:
+    --------
+    retriever : Retriever
+        A retriever object for querying relevant documents.
+    """
+    # Create FAISS index from documents
+    vector_store = FAISS.from_documents(docs, embedding=embeddings)
+
+    # Prepare search kwargs with optional lambda_mult
+    search_kwargs = {'k': k}
+    if lambda_mult is not None:
+        search_kwargs['lambda_mult'] = lambda_mult
+
+    # Return the retriever
+    return vector_store.as_retriever(search_type=search_type, search_kwargs=search_kwargs)