116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
from typing import TypedDict, Any
|
||
import random
|
||
|
||
# LLM setup – use placeholder values if no specific provider is mentioned
|
||
from langchain_openai import ChatOpenAI
|
||
from pydantic import SecretStr
|
||
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b",
|
||
base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
|
||
api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
|
||
temperature=0.2,
|
||
)
|
||
|
||
# ---------- State ----------
|
||
class AgentState(TypedDict):
|
||
task: str
|
||
result: str
|
||
attempts: int
|
||
status: str # pending | success | failed | max_attempts
|
||
error: str | None
|
||
max_attempts: int
|
||
|
||
# ---------- Tool ----------
|
||
def unreliable_tool(task: str) -> str:
|
||
"""Simulate a tool that fails ~30% of the time."""
|
||
if random.random() < 0.3:
|
||
raise ValueError("Tool failure")
|
||
# Very simple evaluation: just return the task string for demo
|
||
return f"Result of '{task}'"
|
||
|
||
# ---------- Nodes ----------
|
||
def execute_task(state: AgentState) -> AgentState:
|
||
try:
|
||
result = unreliable_tool(state["task"])
|
||
state.update(result=result, error=None)
|
||
except Exception as e:
|
||
state.update(result="", error=str(e))
|
||
state.update(status="pending")
|
||
return state
|
||
|
||
def verify_result(state: AgentState) -> AgentState:
|
||
if state["error"]:
|
||
# If tool failed, skip verification
|
||
state.update(status="failed")
|
||
return state
|
||
prompt = f"Task result: {state['result']}. Is this correct? Respond with 'success' or 'failed'."
|
||
verdict_obj = llm.invoke(prompt)
|
||
# Depending on the LLM implementation, the response may be a string or an object with `content`
|
||
if hasattr(verdict_obj, "content"):
|
||
verdict = verdict_obj.content.strip().lower()
|
||
else:
|
||
verdict = str(verdict_obj).strip().lower()
|
||
if "success" in verdict:
|
||
state.update(status="success")
|
||
else:
|
||
state.update(status="failed")
|
||
return state
|
||
|
||
def handle_error(state: AgentState) -> AgentState:
|
||
state["attempts"] += 1
|
||
if state["attempts"] >= state["max_attempts"]:
|
||
state.update(status="max_attempts")
|
||
else:
|
||
state.update(status="pending")
|
||
return state
|
||
|
||
# ---------- Graph ----------
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
|
||
builder = StateGraph(AgentState)
|
||
builder.add_node("execute_task", execute_task)
|
||
builder.add_node("verify_result", verify_result)
|
||
builder.add_node("handle_error", handle_error)
|
||
|
||
builder.set_entry_point("execute_task")
|
||
builder.add_edge("execute_task", "verify_result")
|
||
|
||
def _next(state: AgentState) -> str:
|
||
status = state["status"]
|
||
if status == "success":
|
||
return END
|
||
if status == "failed" and state["attempts"] < state["max_attempts"]:
|
||
return "handle_error"
|
||
return END
|
||
|
||
builder.add_conditional_edges("verify_result", _next)
|
||
builder.add_edge("handle_error", "execute_task")
|
||
|
||
graph = builder.compile(checkpointer=InMemorySaver())
|
||
|
||
# ---------- CLI ----------
|
||
def main():
|
||
task = input("Задача: ").strip()
|
||
if not task:
|
||
print("Нет задачи")
|
||
return
|
||
initial_state: AgentState = {
|
||
"task": task,
|
||
"result": "",
|
||
"attempts": 0,
|
||
"status": "pending",
|
||
"error": None,
|
||
"max_attempts": 5,
|
||
}
|
||
state = graph.invoke(initial_state)
|
||
attempts = state["attempts"] + (1 if state["status"] != "failed" else 0)
|
||
print(f"\nИтог: {state['status']} за {attempts} попытки(й)")
|
||
if state["result"]:
|
||
print(f"Результат: {state['result']}")
|
||
if state["error"]:
|
||
print(f"Ошибка: {state['error']}")
|
||
|
||
if __name__ == "__main__":
|
||
main() |