134 lines
5.1 KiB
Python
134 lines
5.1 KiB
Python
"""
|
||
Main entry point for the LangChain + Qdrant knowledge‑base agent.
|
||
|
||
The script demonstrates three independent usage examples:
|
||
|
||
1. **Simple search** – a single query is sent to the ``search_knowledge_base`` tool.
|
||
2. **Add & search** – a document is added to the collection and then searched.
|
||
3. **Interactive chat** – an agent that can call both tools in a conversational
|
||
setting, using stream mode so that responses appear token‑by‑token.
|
||
|
||
All examples are wrapped in ``if __name__ == "__main__"`` blocks so they run
|
||
only when the module is executed directly.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from typing import Dict, Any
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# LangChain imports – we use only what is required for the examples.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage
|
||
from langchain.agents import create_agent
|
||
from langchain.tools import tool
|
||
|
||
# Import our custom tools
|
||
from .tools import search_knowledge_base, add_to_knowledge_base
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# LLM configuration – the same model is used for all examples.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
|
||
api_key=os.getenv("JOURNAL_MCP_PAT"),
|
||
temperature=0.5,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helper – format a LangChain message for printing.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _format_message(msg: Any) -> str:
|
||
if hasattr(msg, "content") and msg.content:
|
||
return msg.content
|
||
# Fallback to tool call representation
|
||
if hasattr(msg, "tool_calls") and msg.tool_calls:
|
||
tc = msg.tool_calls[0]
|
||
return f"{tc['name']}({tc['args']})"
|
||
return str(msg)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Example 1 – simple search using the tool directly.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def example_simple_search() -> None:
|
||
print("\n=== Example 1: Simple Search ===")
|
||
query = "Python async programming"
|
||
result = search_knowledge_base(query, max_results=3)
|
||
print(f"Query: {query}\nResult:\n{result}")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Example 2 – add a document then search.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def example_add_and_search() -> None:
|
||
print("\n=== Example 2: Add & Search ===")
|
||
content = (
|
||
"Async programming in Python is supported via the asyncio library. "
|
||
"It allows concurrent execution of IO‑bound tasks without threads."
|
||
)
|
||
title = "Python Asyncio"
|
||
add_msg = add_to_knowledge_base(content, title=title)
|
||
print(add_msg)
|
||
|
||
# Now search for a related term.
|
||
query = "asyncio" # short keyword to trigger the newly added doc
|
||
result = search_knowledge_base(query, max_results=2)
|
||
print(f"Search results for '{query}':\n{result}")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Example 3 – interactive chat agent using stream mode.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def example_chat_agent() -> None:
|
||
print("\n=== Example 3: Interactive Chat Agent (stream) ===")
|
||
|
||
# Create an agent that can call our two tools.
|
||
agent = create_agent(
|
||
llm=llm,
|
||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||
system_prompt="You are a helpful assistant with access to a knowledge base. "
|
||
"Use the provided tools to answer user queries.",
|
||
)
|
||
|
||
# Simple chat loop – only one turn for demonstration.
|
||
user_input = "Tell me about async programming in Python."
|
||
print(f"User: {user_input}\n")
|
||
|
||
stream = agent.stream(
|
||
{"messages": [HumanMessage(content=user_input)]},
|
||
stream_mode=["messages", "updates"],
|
||
)
|
||
|
||
step = 1
|
||
for chunk_type, chunk_data in stream:
|
||
if chunk_type == "messages":
|
||
msg, _meta = chunk_data
|
||
# Detect step change – a simple visual separator.
|
||
if _meta.get("langgraph_step") != step:
|
||
step = _meta["langgraph_step"]
|
||
print("\n--- --- --- \n")
|
||
print(_format_message(msg), end="", flush=True)
|
||
elif chunk_type == "updates":
|
||
# When the model finishes a tool call we can show it.
|
||
if chunk_data.get("model"):
|
||
last_msg = chunk_data["model"]["messages"][-1]
|
||
print(_format_message(last_msg))
|
||
|
||
print("\n--- End of conversation ---")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Entry point – run all examples.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
if __name__ == "__main__":
|
||
example_simple_search()
|
||
example_add_and_search()
|
||
example_chat_agent()
|