File size: 4,374 Bytes
dc7b96e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | #!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
from typing import Any
DG_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CONFIG = DG_ROOT / "configs" / "client_profiles" / "langgraph.dg.json"
def load_config(path: Path) -> dict[str, Any]:
return json.loads(path.read_text(encoding="utf-8"))
def import_langgraph() -> tuple[Any, Any, str]:
from langchain_openai import ChatOpenAI
try:
from langchain.agents import create_agent
return ChatOpenAI, create_agent, "langchain.agents.create_agent"
except ImportError:
from langgraph.prebuilt import create_react_agent
return ChatOpenAI, create_react_agent, "langgraph.prebuilt.create_react_agent"
def model_kwargs(config: dict[str, Any]) -> dict[str, Any]:
return {
"model": os.environ.get("LANGGRAPH_MODEL") or config["model"],
"base_url": os.environ.get("OPENAI_BASE_URL") or config["base_url"],
"api_key": os.environ.get("OPENAI_API_KEY") or config["api_key"],
"max_tokens": int(os.environ.get("LANGGRAPH_MAX_TOKENS") or config.get("max_tokens") or 256),
"temperature": float(os.environ.get("LANGGRAPH_TEMPERATURE") or config.get("temperature") or 0.0),
}
def build_agent(config: dict[str, Any]) -> tuple[Any, str]:
ChatOpenAI, factory, factory_name = import_langgraph()
model = ChatOpenAI(**model_kwargs(config))
tools: list[Any] = []
if factory_name == "langchain.agents.create_agent":
return factory(model=model, tools=tools, system_prompt=config.get("system_prompt")), factory_name
return factory(model, tools, prompt=config.get("system_prompt")), factory_name
def last_message_content(result: Any) -> str:
if isinstance(result, dict):
messages = result.get("messages") or []
if messages:
last = messages[-1]
return str(getattr(last, "content", last.get("content") if isinstance(last, dict) else last))
return str(result)
def run_task(args: argparse.Namespace, config: dict[str, Any]) -> int:
agent, factory_name = build_agent(config)
result = agent.invoke({"messages": [{"role": "user", "content": args.task}]})
if args.json:
print(json.dumps({"status": "success", "agent_factory": factory_name, "result": last_message_content(result)}, ensure_ascii=False, indent=2))
else:
print(last_message_content(result))
return 0
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run LangGraph/LangChain agent with the local DiffusionGemma profile.")
parser.add_argument("--repo", default=".", help="Target repo, used as working directory")
parser.add_argument("--config", default=str(DEFAULT_CONFIG))
parser.add_argument("--task", default="")
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--smoke-import", action="store_true")
parser.add_argument("--json", action="store_true")
return parser.parse_args()
def main() -> int:
args = parse_args()
repo = Path(args.repo).resolve()
config_path = Path(args.config).resolve()
config = load_config(config_path)
if args.smoke_import:
ChatOpenAI, factory, factory_name = import_langgraph()
print("langgraph import ok")
print(ChatOpenAI.__name__)
print(factory.__name__)
print(factory_name)
return 0
if args.dry_run:
_, _, factory_name = import_langgraph()
data = {
"repo": str(repo),
"config": str(config_path),
"agent_factory": factory_name,
"configured_agent_factory": config["agent_factory"],
"fallback_agent_factory": config["fallback_agent_factory"],
"model_class": config["model_class"],
"model_kwargs": model_kwargs(config),
"command": f"scripts/dg_agent.sh langgraph -- --repo {repo} --task '...'",
}
print(json.dumps(data, ensure_ascii=False, indent=2) if args.json else "\n".join(f"{k}: {v}" for k, v in data.items()))
return 0
if not args.task:
print("--task is required unless --dry-run or --smoke-import is used", flush=True)
return 2
os.chdir(repo)
return run_task(args, config)
if __name__ == "__main__":
raise SystemExit(main())
|