Add main.py

This commit is contained in:
2026-06-04 15:55:36 +00:00
parent f9f8224e18
commit dd9d3c726b
+132
View File
@@ -0,0 +1,132 @@
"""
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 singleturn 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
# Thirdparty 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]Singleturn 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 80line requirement.
It demonstrates synchronous singleturn, 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 objectoriented 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]")
"""