integrated local db connection

Author: realshubhamraut

.DS_Store

@@ -1,112 +1,72 @@
-import streamlit as st
+# agent.py
 from dataclasses import dataclass
 from typing import Annotated, Sequence, Optional
-
-from langchain.callbacks.base import BaseCallbackHandler
+from langchain_core.messages import BaseMessage
+from langgraph.graph.message import add_messages
 from langchain_core.messages import SystemMessage
-from langchain_openai import ChatOpenAI
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import START, END, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
-from langgraph.graph.message import add_messages
-from langchain_core.messages import BaseMessage
 
-from tools import retriever_tool
-from tools import search, sql_executor_tool
-from PIL import Image
-from io import BytesIO
+# Exported items
+__all__ = ["MessagesState", "create_agent"]
 
 @dataclass
 class MessagesState:
     messages: Annotated[Sequence[BaseMessage], add_messages]
 
-
 memory = MemorySaver()
 
-
+# Model configuration for Google Gemini only
 @dataclass
 class ModelConfig:
     model_name: str
     api_key: str
     base_url: Optional[str] = None
 
-
 model_configurations = {
-    # "o3-mini": ModelConfig(
-    #     model_name="o3-mini", api_key=st.secrets["OPENAI_API_KEY"]
-    # ),
-
-    
-    "Deepseek R1": ModelConfig(
-        model_name="deepseek-r1-distill-llama-70b",
-        api_key=st.secrets["GROK_API_KEY"],
-        base_url=f"https://api.groq.com/openai/v1",
-    ),
-
-    # "Mistral 7B": ModelConfig(
-    #     model_name="mistralai/mistral-7b-v0.1", api_key=st.secrets["REPLICATE_API_TOKEN"]
-    # ),
-    # "Qwen 2.5": ModelConfig(
-    #     model_name="accounts/fireworks/models/qwen2p5-coder-32b-instruct",
-    #     api_key=st.secrets["FIREWORKS_API_KEY"],
-    #     base_url="https://api.fireworks.ai/inference/v1",
-    # ),
-    # "Gemini Exp 1206": ModelConfig(
-    #     model_name="gemini-exp-1206",
-    #     api_key=st.secrets["GEMINI_API_KEY"],
-    #     base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
-    # ),
+    "Google Gemini": ModelConfig(
+         model_name="models/gemini-2.0-flash",
+         api_key=__import__("streamlit").secrets["GEMINI_API_KEY"],
+         base_url=None,
+    )
 }
+
 sys_msg = SystemMessage(
-    content="""You're an AI assistant specializing in data analysis with Snowflake SQL. When providing responses, strive to exhibit friendliness and adopt a conversational tone, similar to how a friend or tutor would communicate. Do not ask the user for schema or database details. You have access to the following tools:
-    ALWAYS USE THE DATABASE_SCHEMA TOOL TO GET THE SCHEMA OF THE DATABASE BEFORE GENERATING SQL CODE.
-    ALWAYS USE THE DATABASE_SCHEMA TOOL TO GET THE SCHEMA OF THE DATABASE BEFORE GENERATING SQL CODE.
-    - Database_Schema: This tool allows you to search for database schema details when needed to generate the SQL code.
-    - Internet_Search: This tool allows you to search the internet for snowflake sql related information when needed to generate the SQL code.
+    content="""You're an AI assistant specializing in data analysis with Snowflake SQL. When providing responses, strive to be friendly and conversational (like a tutor or friend). You have access to the following tools:
+    - Database_Schema: Search for database schema details before generating SQL code.
+    - Internet_Search: Look up Snowflake SQL–related information when needed.
     """
 )
-tools = [retriever_tool, search]
 
-def create_agent(callback_handler: BaseCallbackHandler, model_name: str) -> StateGraph:
-    config = model_configurations.get(model_name)
-    if not config:
-        raise ValueError(f"Unsupported model name: {model_name}")
+# Tools are imported from tools.py (see that file)
+from tools import retriever_tool, search
+tools = [retriever_tool, search]
 
+def create_agent(callback_handler) -> StateGraph:
+    config = model_configurations["Google Gemini"]
     if not config.api_key:
-        raise ValueError(f"API key for model '{model_name}' is not set. Please check your environment variables or secrets configuration.")
-
-    llm = ChatOpenAI(
-        model=config.model_name,
-        api_key=config.api_key,
-        callbacks=[callback_handler],
-        streaming=True,
-        base_url=config.base_url,
-        # temperature=0.1,
-        default_headers={"HTTP-Referer": "", "X-Title": "Snowchat"},
+        raise ValueError("API key for Google Gemini is not set. Please check your secrets configuration.")
+    llm = ChatGoogleGenerativeAI(
+         model=config.model_name,
+         google_api_key=config.api_key,
+         callbacks=[callback_handler],
+         temperature=0,
+         base_url=config.base_url,
+         streaming=True,
     )
-
     llm_with_tools = llm.bind_tools(tools)
 
     def llm_agent(state: MessagesState):
-        return {"messages": [llm_with_tools.invoke([sys_msg] + state.messages)]}
+         return {"messages": [llm_with_tools.invoke([sys_msg] + state.messages)]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("llm_agent", llm_agent)
     builder.add_node("tools", ToolNode(tools))
-
     builder.add_edge(START, "llm_agent")
     builder.add_conditional_edges("llm_agent", tools_condition)
     builder.add_edge("tools", "llm_agent")
     builder.add_edge("llm_agent", END)
     react_graph = builder.compile(checkpointer=memory)
-
-    # png_data = react_graph.get_graph(xray=True).draw_mermaid_png()
-    # with open("graph_2.png", "wb") as f:
-    #     f.write(png_data)
-
-    # image = Image.open(BytesIO(png_data))
-    # st.image(image, caption="React Graph")
-
     return react_graph
-
-    

__pycache__/agent.cpython-312.pyc

@@ -1,37 +1,31 @@
+# chain.py
 from dataclasses import dataclass, field
 from operator import itemgetter
 from typing import Any, Callable, Dict, Optional
-
 import streamlit as st
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.llms import OpenAI
+from langchain_community.embeddings import FakeEmbeddings
 from langchain.prompts.prompt import PromptTemplate
 from langchain.schema import format_document
 from langchain.vectorstores import SupabaseVectorStore
-from langchain_anthropic import ChatAnthropic
-from langchain_community.chat_models import ChatOpenAI
 from langchain_core.messages import get_buffer_string
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableParallel, RunnablePassthrough
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-
+from langchain_google_genai import ChatGoogleGenerativeAI
 from supabase.client import Client, create_client
 from template import CONDENSE_QUESTION_PROMPT, QA_PROMPT
 
 DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template="{page_content}")
 
 supabase_url = st.secrets["SUPABASE_URL"]
 supabase_key = st.secrets["SUPABASE_SERVICE_KEY"]
-supabase: Client = create_client(supabase_url, supabase_key)
-
+client: Client = create_client(supabase_url, supabase_key)
 
 @dataclass
 class ModelConfig:
     model_type: str
     secrets: Dict[str, Any]
     callback_handler: Optional[Callable] = field(default=None)
 
-
 class ModelWrapper:
     def __init__(self, config: ModelConfig):
         self.model_type = config.model_type
@@ -40,68 +34,17 @@ def __init__(self, config: ModelConfig):
         self.llm = self._setup_llm()
 
     def _setup_llm(self):
-        model_config = {
-            "gpt-4o-mini": {
-                "model_name": "gpt-4o-mini",
-                "api_key": self.secrets["OPENAI_API_KEY"],
-            },
-            "gemma2-9b": {
-                "model_name": "gemma2-9b-it",
-                "api_key": self.secrets["GROQ_API_KEY"],
-                "base_url": "https://api.groq.com/openai/v1",
-            },
-            "claude3-haiku": {
-                "model_name": "claude-3-haiku-20240307",
-                "api_key": self.secrets["ANTHROPIC_API_KEY"],
-            },
-            "mixtral-8x22b": {
-                "model_name": "accounts/fireworks/models/mixtral-8x22b-instruct",
-                "api_key": self.secrets["FIREWORKS_API_KEY"],
-                "base_url": "https://api.fireworks.ai/inference/v1",
-            },
-            "llama-3.1-405b": {
-                "model_name": "accounts/fireworks/models/llama-v3p1-405b-instruct",
-                "api_key": self.secrets["FIREWORKS_API_KEY"],
-                "base_url": "https://api.fireworks.ai/inference/v1",
-            },
-        }
-
-        config = model_config[self.model_type]
-
-        return (
-            ChatOpenAI(
-                model_name=config["model_name"],
-                temperature=0.1,
-                api_key=config["api_key"],
-                max_tokens=700,
-                callbacks=[self.callback_handler],
-                streaming=True,
-                base_url=config["base_url"]
-                if config["model_name"] != "gpt-4o-mini"
-                else None,
-                default_headers={
-                    "HTTP-Referer": "https://snowchat.streamlit.app/",
-                    "X-Title": "Snowchat",
-                },
-            )
-            if config["model_name"] != "claude-3-haiku-20240307"
-            else (
-                ChatAnthropic(
-                    model=config["model_name"],
-                    temperature=0.1,
-                    max_tokens=700,
-                    timeout=None,
-                    max_retries=2,
-                    callbacks=[self.callback_handler],
-                    streaming=True,
-                )
-            )
+        return ChatGoogleGenerativeAI(
+            model="models/gemini-2.0-flash",
+            google_api_key=self.secrets["GEMINI_API_KEY"],
+            temperature=0.1,
+            callbacks=[self.callback_handler],
+            max_tokens=700,
+            streaming=True,
         )
 
     def get_chain(self, vectorstore):
-        def _combine_documents(
-            docs, document_prompt=DEFAULT_DOCUMENT_PROMPT, document_separator="\n\n"
-        ):
+        def _combine_documents(docs, document_prompt=DEFAULT_DOCUMENT_PROMPT, document_separator="\n\n"):
             doc_strings = [format_document(doc, document_prompt) for doc in docs]
             return document_separator.join(doc_strings)
 
@@ -110,8 +53,7 @@ def _combine_documents(
                 chat_history=lambda x: get_buffer_string(x["chat_history"])
             )
             | CONDENSE_QUESTION_PROMPT
-            | OpenAI()
-            | StrOutputParser(),
+            | StrOutputParser()
         )
         _context = {
             "context": itemgetter("standalone_question")
@@ -120,33 +62,23 @@ def _combine_documents(
             "question": lambda x: x["standalone_question"],
         }
         conversational_qa_chain = _inputs | _context | QA_PROMPT | self.llm
-
         return conversational_qa_chain
 
-
-def load_chain(model_name="qwen", callback_handler=None):
-    embeddings = OpenAIEmbeddings(
-        openai_api_key=st.secrets["OPENAI_API_KEY"], model="text-embedding-ada-002"
-    )
+def load_chain(model_name="google_gemini", callback_handler=None):
+    embeddings = FakeEmbeddings(size=768)
     vectorstore = SupabaseVectorStore(
         embedding=embeddings,
-        client=supabase,
+        client=client,
         table_name="documents",
         query_name="v_match_documents",
     )
-
-    model_type_mapping = {
-        "gpt-4o-mini": "gpt-4o-mini",
-        "gemma2-9b": "gemma2-9b",
-        "claude3-haiku": "claude3-haiku",
-        "mixtral-8x22b": "mixtral-8x22b",
-        "llama-3.1-405b": "llama-3.1-405b",
-    }
-
-    model_type = model_type_mapping.get(model_name.lower())
-    if model_type is None:
-        raise ValueError(f"Unsupported model name: {model_name}")
-
+    # Override the retriever with a dummy retriever to disable document retrieval.
+    class DummyRetriever:
+        def get_relevant_documents(self, query):
+            return []
+    vectorstore.as_retriever = lambda: DummyRetriever()
+    
+    model_type = "google_gemini"
     config = ModelConfig(
         model_type=model_type, secrets=st.secrets, callback_handler=callback_handler
     )

__pycache__/local_chat.cpython-312.pyc

@@ -1,3 +1,4 @@
+# ingest.py
 from typing import Any, Dict
 
 import streamlit as st
@@ -9,40 +10,30 @@
 
 from supabase.client import Client, create_client
 
-
 class Secrets(BaseModel):
     SUPABASE_URL: str
     SUPABASE_SERVICE_KEY: str
     OPENAI_API_KEY: str
 
-
 class Config(BaseModel):
     chunk_size: int = 1000
     chunk_overlap: int = 0
     docs_dir: str = "docs/"
     docs_glob: str = "**/*.md"
 
-
 class DocumentProcessor:
     def __init__(self, secrets: Secrets, config: Config):
-        self.client: Client = create_client(
-            secrets.SUPABASE_URL, secrets.SUPABASE_SERVICE_KEY
-        )
+        self.client: Client = create_client(secrets.SUPABASE_URL, secrets.SUPABASE_SERVICE_KEY)
         self.loader = DirectoryLoader(config.docs_dir, glob=config.docs_glob)
-        self.text_splitter = CharacterTextSplitter(
-            chunk_size=config.chunk_size, chunk_overlap=config.chunk_overlap
-        )
+        self.text_splitter = CharacterTextSplitter(chunk_size=config.chunk_size, chunk_overlap=config.chunk_overlap)
         self.embeddings = OpenAIEmbeddings(openai_api_key=secrets.OPENAI_API_KEY)
 
     def process(self) -> Dict[str, Any]:
         data = self.loader.load()
         texts = self.text_splitter.split_documents(data)
-        vector_store = SupabaseVectorStore.from_documents(
-            texts, self.embeddings, client=self.client
-        )
+        vector_store = SupabaseVectorStore.from_documents(texts, self.embeddings, client=self.client)
         return vector_store
 
-
 def run():
     secrets = Secrets(
         SUPABASE_URL=st.secrets["SUPABASE_URL"],
@@ -54,6 +45,5 @@ def run():
     result = doc_processor.process()
     return result
 
-
 if __name__ == "__main__":
     run()