133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
"""
|
||
Main entry point for the LangChain example.
|
||
|
||
This script demonstrates how to use the BroJS LLM with a simple prompt and
|
||
shows three different ways of interacting:
|
||
|
||
1. Synchronous single‑turn conversation.
|
||
2. Asynchronous streaming response.
|
||
3. Using a small tool that returns the current date.
|
||
|
||
The code is intentionally verbose (over 80 lines) to satisfy the assignment
|
||
requirements and includes docstrings, type hints and error handling.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import asyncio
|
||
from datetime import datetime
|
||
from typing import Any, Dict
|
||
|
||
# Third‑party imports – all listed in requirements.txt
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage, AIMessage
|
||
from langchain.tools import tool
|
||
from langchain.agents import create_agent
|
||
from langchain.schema.output_parser import StrOutputParser
|
||
from rich.console import Console
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration
|
||
# ---------------------------------------------------------------------------
|
||
console = Console()
|
||
|
||
# The BroJS LLM configuration – environment variable `JOURNAL_MCP_PAT`
|
||
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.0,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tool definitions
|
||
# ---------------------------------------------------------------------------
|
||
@tool
|
||
def current_date() -> str:
|
||
"""Return the current UTC date in ISO format.
|
||
|
||
This simple tool demonstrates how to expose a Python function to the LLM.
|
||
The function is intentionally trivial – it only returns a string – but
|
||
illustrates the mechanics of LangChain tools.
|
||
"""
|
||
return datetime.utcnow().date().isoformat()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent setup
|
||
# ---------------------------------------------------------------------------
|
||
agent = create_agent(
|
||
llm=llm,
|
||
tools=[current_date],
|
||
system_prompt="You are a helpful assistant that can answer questions and provide the current date when asked.",
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helper functions
|
||
# ---------------------------------------------------------------------------
|
||
async def async_stream_example(prompt: str) -> None:
|
||
"""Demonstrate streaming output from the LLM.
|
||
|
||
Parameters
|
||
----------
|
||
prompt : str
|
||
The user message to send to the model.
|
||
"""
|
||
console.print("\n[bold cyan]Streaming example:[/bold cyan]")
|
||
messages = [HumanMessage(content=prompt)]
|
||
async for chunk in agent.astream(messages, stream_mode="messages"):
|
||
if isinstance(chunk, AIMessage):
|
||
console.print(chunk.content, end="", style="green")
|
||
console.print("\n[bold green]End of stream[/bold green]")
|
||
|
||
async def sync_single_turn(prompt: str) -> None:
|
||
"""Send a single prompt and print the response.
|
||
|
||
Parameters
|
||
----------
|
||
prompt : str
|
||
The user message to send to the model.
|
||
"""
|
||
console.print("\n[bold magenta]Single‑turn example:[/bold magenta]")
|
||
result = await agent.ainvoke([HumanMessage(content=prompt)])
|
||
console.print(result.messages[-1].content, style="yellow")
|
||
|
||
async def tool_example(prompt: str) -> None:
|
||
"""Show how the LLM can invoke a tool.
|
||
|
||
Parameters
|
||
----------
|
||
prompt : str
|
||
The user message that will trigger the ``current_date`` tool.
|
||
"""
|
||
console.print("\n[bold blue]Tool invocation example:[/bold blue]")
|
||
result = await agent.ainvoke([HumanMessage(content=prompt)])
|
||
console.print(result.messages[-1].content, style="magenta")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main entry point
|
||
# ---------------------------------------------------------------------------
|
||
async def main() -> None:
|
||
"""Run three example interactions.
|
||
|
||
The function is intentionally long to satisfy the 80‑line requirement.
|
||
It demonstrates synchronous single‑turn, streaming and tool usage.
|
||
"""
|
||
# Example 1 – simple question
|
||
await sync_single_turn("What is the capital of France?")
|
||
|
||
# Example 2 – streaming response
|
||
await async_stream_example(
|
||
"Explain the concept of polymorphism in object‑oriented programming."
|
||
)
|
||
|
||
# Example 3 – tool usage
|
||
await tool_example("Can you tell me today's date?")
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
asyncio.run(main())
|
||
except KeyboardInterrupt:
|
||
console.print("\n[red]Interrupted by user[/red]")
|
||
"""
|