Move to Langgraph Agents - Its time

Makefile

@@ -0,0 +1,57 @@
+.PHONY: all format lint lint_diff format_diff lint_package lint_tests spell_check spell_fix help lint-fix
+
+# Define a variable for Python and notebook files.
+PYTHON_FILES=src/
+MYPY_CACHE=.mypy_cache
+
+######################
+# LINTING AND FORMATTING
+######################
+
+lint format: PYTHON_FILES=.
+lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --diff-filter=d main | grep -E '\.py$$|\.ipynb$$')
+lint_package: PYTHON_FILES=src
+lint_tests: PYTHON_FILES=tests
+lint_tests: MYPY_CACHE=.mypy_cache_test
+
+lint lint_diff lint_package lint_tests:
+	python -m ruff check .
+	[ "$(PYTHON_FILES)" = "" ] || python -m ruff format $(PYTHON_FILES) --diff
+	[ "$(PYTHON_FILES)" = "" ] || python -m ruff check --select I,F401 --fix $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || python -m mypy --strict $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) && python -m mypy --strict $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
+
+format format_diff:
+	ruff format $(PYTHON_FILES)
+	ruff check --fix $(PYTHON_FILES)
+
+spell_check:
+	codespell --toml pyproject.toml
+
+spell_fix:
+	codespell --toml pyproject.toml -w
+
+######################
+# RUN ALL
+######################
+
+all: format lint spell_check
+
+######################
+# HELP
+######################
+
+help:
+	@echo '----'
+	@echo 'format                       - run code formatters'
+	@echo 'lint                         - run linters'
+	@echo 'spell_check                  - run spell check'
+	@echo 'all                          - run all tasks'
+	@echo 'lint-fix                     - run lint and fix issues'
+
+######################
+# LINT-FIX TARGET
+######################
+
+lint-fix: format lint
+	@echo "Linting and fixing completed successfully."

README.md

@@ -15,9 +15,11 @@
 
 ## Supported LLM's
 
-- GPT-3.5-turbo-0125
-- CodeLlama-70B
-- Mistral Medium
+- GPT-4o
+- Gemini Flash 1.5 8B
+- Claude 3 Haiku
+- Llama 3.2 3B
+- Llama 3.1 405B
 
 #
 
@@ -27,11 +29,12 @@ https://github.com/kaarthik108/snowChat/assets/53030784/24105e23-69d3-4676-b6d6-
 
 ## 🌟 Features
 
-- **Conversational AI**: Harnesses ChatGPT to translate natural language into precise SQL queries.
+- **Conversational AI**: Use ChatGPT and other models to translate natural language into precise SQL queries.
 - **Conversational Memory**: Retains context for interactive, dynamic responses.
 - **Snowflake Integration**: Offers seamless, real-time data insights straight from your Snowflake database.
 - **Self-healing SQL**: Proactively suggests solutions for SQL errors, streamlining data access.
 - **Interactive User Interface**: Transforms data querying into an engaging conversation, complete with a chat reset option.
+- **Agent-based Architecture**: Utilizes an agent to manage interactions and tool usage.
 
 ## 🛠️ Installation
 
@@ -42,7 +45,9 @@ https://github.com/kaarthik108/snowChat/assets/53030784/24105e23-69d3-4676-b6d6-
    cd snowchat
    pip install -r requirements.txt
 
-3. Set up your `OPENAI_API_KEY`, `ACCOUNT`, `USER_NAME`, `PASSWORD`, `ROLE`, `DATABASE`, `SCHEMA`, `WAREHOUSE`, `SUPABASE_URL` , `SUPABASE_SERVICE_KEY` and `REPLICATE_API_TOKEN` in project directory `secrets.toml`.
+3. Set up your `OPENAI_API_KEY`, `ACCOUNT`, `USER_NAME`, `PASSWORD`, `ROLE`, `DATABASE`, `SCHEMA`, `WAREHOUSE`, `SUPABASE_URL` , `SUPABASE_SERVICE_KEY`, `SUPABASE_STORAGE_URL`,`CLOUDFLARE_ACCOUNT_ID`, `CLOUDFLARE_NAMESPACE_ID`,
+   `CLOUDFLARE_API_TOKEN` in project directory `secrets.toml`.
+   Cloudflare is used here for caching Snowflake responses in KV.
 
 4. Make you're schemas and store them in docs folder that matches you're database.
 
@@ -53,12 +58,6 @@ https://github.com/kaarthik108/snowChat/assets/53030784/24105e23-69d3-4676-b6d6-
 7. Run the Streamlit app to start chatting:
    streamlit run main.py
 
-## 🚀 Additional Enhancements
-
-1. **Platform Integration**: Connect snowChat with popular communication platforms like Slack or Discord for seamless interaction.
-2. **Voice Integration**: Implement voice recognition and text-to-speech functionality to make the chatbot more interactive and user-friendly.
-3. **Advanced Analytics**: Integrate with popular data visualization libraries like Plotly or Matplotlib to generate interactive visualizations based on the user's queries (AutoGPT).
-
 ## Star History
 
 [![Star History Chart](https://api.star-history.com/svg?repos=kaarthik108/snowChat&type=Date)]

agent.py

@@ -0,0 +1,115 @@
+import os
+import streamlit as st
+from dataclasses import dataclass
+from typing import Annotated, Sequence, Optional
+
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain_anthropic import ChatAnthropic
+from langchain_core.messages import SystemMessage
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import START, 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
+
+@dataclass
+class MessagesState:
+    messages: Annotated[Sequence[BaseMessage], add_messages]
+
+
+memory = MemorySaver()
+
+
+@dataclass
+class ModelConfig:
+    model_name: str
+    api_key: str
+    base_url: Optional[str] = None
+
+
+model_configurations = {
+    "gpt-4o": ModelConfig(
+        model_name="gpt-4o", api_key=st.secrets["OPENAI_API_KEY"]
+    ),
+    "Gemini Flash 1.5 8B": ModelConfig(
+        model_name="google/gemini-flash-1.5-8b",
+        api_key=st.secrets["OPENROUTER_API_KEY"],
+        base_url="https://openrouter.ai/api/v1",
+    ),
+    "claude3-haiku": ModelConfig(
+        model_name="claude-3-haiku-20240307", api_key=st.secrets["ANTHROPIC_API_KEY"]
+    ),
+    "llama-3.2-3b": ModelConfig(
+        model_name="accounts/fireworks/models/llama-v3p2-3b-instruct",
+        api_key=st.secrets["FIREWORKS_API_KEY"],
+        base_url="https://api.fireworks.ai/inference/v1",
+    ),
+    "llama-3.1-405b": ModelConfig(
+        model_name="accounts/fireworks/models/llama-v3p1-405b-instruct",
+        api_key=st.secrets["FIREWORKS_API_KEY"],
+        base_url="https://api.fireworks.ai/inference/v1",
+    ),
+}
+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:
+    - 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.
+    - Snowflake_SQL_Executor: This tool allows you to execute snowflake sql queries when needed to generate the SQL code. You only have read access to the database, do not modify the database in any way.
+    """
+)
+tools = [retriever_tool, search, sql_executor_tool]
+
+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}")
+
+    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.01,
+        )
+        if config.model_name != "claude-3-haiku-20240307"
+        else ChatAnthropic(
+            model=config.model_name,
+            api_key=config.api_key,
+            callbacks=[callback_handler],
+            streaming=True,
+        )
+    )
+
+    llm_with_tools = llm.bind_tools(tools)
+
+    def llm_agent(state: MessagesState):
+        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")
+
+    react_graph = builder.compile(checkpointer=memory)
+
+    # png_data = react_graph.get_graph(xray=True).draw_mermaid_png()
+    # with open("graph.png", "wb") as f:
+    #     f.write(png_data)
+
+    # image = Image.open(BytesIO(png_data))
+    # st.image(image, caption="React Graph")
+
+    return react_graph