Temporarily moved graph to gradio file.

app.py

@@ -1,15 +1,22 @@
 from typing import Any
 
 import gradio as gr
+from langchain.chat_models import init_chat_model
+from langchain_core.tools import tool
 from langchain_openai import OpenAIEmbeddings
 from langchain_qdrant import QdrantVectorStore
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import MessagesState, StateGraph, END
+from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.prebuilt import ToolNode
 import structlog
 
 from qdrant_client import QdrantClient
 from qdrant_client.http.models import Distance, VectorParams, SparseVectorParams
 
 import logging_config as _
-from conversation.main import graph
+# from conversation.main import graph
+from conversation.generate import generate
 from ingestion.main import ingest_document
 
 from config import app_settings
@@ -22,6 +29,12 @@ embeddings = OpenAIEmbeddings(
     api_key=app_settings.llm_api_key
 )
 
+llm = init_chat_model(
+    app_settings.llm_model,
+    model_provider="openai",
+    api_key=app_settings.llm_api_key
+)
+
 client = QdrantClient(app_settings.vector_db_url)
 if not client.collection_exists(app_settings.vector_db_collection_name):
     client.create_collection(
@@ -29,13 +42,55 @@ if not client.collection_exists(app_settings.vector_db_collection_name):
         vectors_config=VectorParams(size=app_settings.embedding_size, distance=Distance.COSINE),
         sparse_vectors_config={'langchain-sparse': SparseVectorParams(index=None, modifier=None)}
     )
+# TODO: move to LLM files later
 vector_store = QdrantVectorStore(
     client=client,
     collection_name=app_settings.vector_db_collection_name,
     embedding=embeddings,
 )
 
+# ------
+# Move to `conversation/main`` later
+@tool(response_format="content_and_artifact")
+def retrieve(query: str):
+    """Retrieve information related to a query."""
+    retrieved_docs = vector_store.similarity_search(query, k=2)
+    serialized = "\n\n".join(
+        (f"Source: {doc.metadata}\n" f"Content: {doc.page_content}")
+        for doc in retrieved_docs
+    )
+    return serialized, retrieved_docs
+
+
+def query_or_respond(state: MessagesState):
+    """Generate tool call for retrieval or respond."""
+    llm_with_tools = llm.bind_tools([retrieve])
+    response = llm_with_tools.invoke(state["messages"])
+    # MessagesState appends messages to state instead of overwriting
+    return {"messages": [response]}
+
+
+graph_builder = StateGraph(MessagesState)
+tools = ToolNode([retrieve])
+memory = MemorySaver()
+
+graph_builder.add_node(query_or_respond)
+graph_builder.add_node(tools)
+graph_builder.add_node(generate)
+
+graph_builder.set_entry_point("query_or_respond")
+
+graph_builder.add_conditional_edges(
+    "query_or_respond",
+    tools_condition,
+    {END: END, "tools": "tools"},
+)
+
+graph_builder.add_edge("tools", "generate")
+graph_builder.add_edge("generate", END)
 
+graph = graph_builder.compile(checkpointer=memory)
+# -----
 
 with open("static/style.css", "r") as f:
     css = f.read()

conversation/main.py

@@ -1,79 +1,79 @@
-from langchain.chat_models import init_chat_model
-from langchain_core.tools import tool
-from langchain_openai import OpenAIEmbeddings
-from langchain_qdrant import QdrantVectorStore
-from langgraph.checkpoint.memory import MemorySaver
-from langgraph.graph import MessagesState, StateGraph, END
-from langgraph.prebuilt import ToolNode, tools_condition
-from langgraph.prebuilt import ToolNode
-from qdrant_client import QdrantClient
-from qdrant_client.http.models import Distance, VectorParams, SparseVectorParams
-
-from config import app_settings
-from conversation.generate import generate
-
-
-llm = init_chat_model(
-    app_settings.llm_model,
-    model_provider="openai",
-    api_key=app_settings.llm_api_key
-)
-
-embeddings = OpenAIEmbeddings(
-    model=app_settings.embedding_model,
-    api_key=app_settings.llm_api_key
-)
-
-client = QdrantClient(app_settings.vector_db_url)
-if not client.collection_exists(app_settings.vector_db_collection_name):
-    client.create_collection(
-        collection_name=app_settings.vector_db_collection_name,
-        vectors_config=VectorParams(size=app_settings.embedding_size, distance=Distance.COSINE),
-        sparse_vectors_config={'langchain-sparse': SparseVectorParams(index=None, modifier=None)}
-    )
-
-vector_store = QdrantVectorStore(
-    client=client,
-    collection_name=app_settings.vector_db_collection_name,
-    embedding=embeddings,
-)
-
-@tool(response_format="content_and_artifact")
-def retrieve(query: str):
-    """Retrieve information related to a query."""
-    retrieved_docs = vector_store.similarity_search(query, k=2)
-    serialized = "\n\n".join(
-        (f"Source: {doc.metadata}\n" f"Content: {doc.page_content}")
-        for doc in retrieved_docs
-    )
-    return serialized, retrieved_docs
-
-
-def query_or_respond(state: MessagesState):
-    """Generate tool call for retrieval or respond."""
-    llm_with_tools = llm.bind_tools([retrieve])
-    response = llm_with_tools.invoke(state["messages"])
-    # MessagesState appends messages to state instead of overwriting
-    return {"messages": [response]}
-
-
-graph_builder = StateGraph(MessagesState)
-tools = ToolNode([retrieve])
-memory = MemorySaver()
-
-graph_builder.add_node(query_or_respond)
-graph_builder.add_node(tools)
-graph_builder.add_node(generate)
-
-graph_builder.set_entry_point("query_or_respond")
-
-graph_builder.add_conditional_edges(
-    "query_or_respond",
-    tools_condition,
-    {END: END, "tools": "tools"},
-)
-
-graph_builder.add_edge("tools", "generate")
-graph_builder.add_edge("generate", END)
-
-graph = graph_builder.compile(checkpointer=memory)
+# from langchain.chat_models import init_chat_model
+# from langchain_core.tools import tool
+# from langchain_openai import OpenAIEmbeddings
+# from langchain_qdrant import QdrantVectorStore
+# from langgraph.checkpoint.memory import MemorySaver
+# from langgraph.graph import MessagesState, StateGraph, END
+# from langgraph.prebuilt import ToolNode, tools_condition
+# from langgraph.prebuilt import ToolNode
+# from qdrant_client import QdrantClient
+# from qdrant_client.http.models import Distance, VectorParams, SparseVectorParams
+
+# from config import app_settings
+# from conversation.generate import generate
+
+
+# llm = init_chat_model(
+#     app_settings.llm_model,
+#     model_provider="openai",
+#     api_key=app_settings.llm_api_key
+# )
+
+# embeddings = OpenAIEmbeddings(
+#     model=app_settings.embedding_model,
+#     api_key=app_settings.llm_api_key
+# )
+
+# client = QdrantClient(app_settings.vector_db_url)
+# if not client.collection_exists(app_settings.vector_db_collection_name):
+#     client.create_collection(
+#         collection_name=app_settings.vector_db_collection_name,
+#         vectors_config=VectorParams(size=app_settings.embedding_size, distance=Distance.COSINE),
+#         sparse_vectors_config={'langchain-sparse': SparseVectorParams(index=None, modifier=None)}
+#     )
+
+# vector_store = QdrantVectorStore(
+#     client=client,
+#     collection_name=app_settings.vector_db_collection_name,
+#     embedding=embeddings,
+# )
+
+# @tool(response_format="content_and_artifact")
+# def retrieve(query: str):
+#     """Retrieve information related to a query."""
+#     retrieved_docs = vector_store.similarity_search(query, k=2)
+#     serialized = "\n\n".join(
+#         (f"Source: {doc.metadata}\n" f"Content: {doc.page_content}")
+#         for doc in retrieved_docs
+#     )
+#     return serialized, retrieved_docs
+
+
+# def query_or_respond(state: MessagesState):
+#     """Generate tool call for retrieval or respond."""
+#     llm_with_tools = llm.bind_tools([retrieve])
+#     response = llm_with_tools.invoke(state["messages"])
+#     # MessagesState appends messages to state instead of overwriting
+#     return {"messages": [response]}
+
+
+# graph_builder = StateGraph(MessagesState)
+# tools = ToolNode([retrieve])
+# memory = MemorySaver()
+
+# graph_builder.add_node(query_or_respond)
+# graph_builder.add_node(tools)
+# graph_builder.add_node(generate)
+
+# graph_builder.set_entry_point("query_or_respond")
+
+# graph_builder.add_conditional_edges(
+#     "query_or_respond",
+#     tools_condition,
+#     {END: END, "tools": "tools"},
+# )
+
+# graph_builder.add_edge("tools", "generate")
+# graph_builder.add_edge("generate", END)
+
+# graph = graph_builder.compile(checkpointer=memory)