Summer Shen
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
from llama_agent_skills.agent import create_agent, run_agent
|
|
from llama_agent_skills.config import Config
|
|
from llama_agent_skills.registry import SkillRegistry
|
|
|
|
|
|
def _parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="LlamaIndex Agent with SKILL support")
|
|
parser.add_argument(
|
|
"--skill-dir",
|
|
type=Path,
|
|
action="append",
|
|
default=None,
|
|
help="Path to a skill directory (can be specified multiple times)",
|
|
)
|
|
parser.add_argument("--model", type=str, default=None, help="LLM model name")
|
|
return parser.parse_args()
|
|
|
|
|
|
async def _run_interactive(config: Config, registry: SkillRegistry) -> None:
|
|
skills = registry.list_skills()
|
|
agent = create_agent(
|
|
skills=skills,
|
|
model=config.llm_model,
|
|
api_key=config.llm_api_key,
|
|
base_url=config.llm_base_url or None,
|
|
base_system_prompt=config.base_system_prompt,
|
|
)
|
|
|
|
loaded_names = [s.name for s in skills]
|
|
print(f"Agent ready. Loaded {len(skills)} skill(s): {loaded_names}")
|
|
print("Type 'quit' or 'exit' to stop.\n")
|
|
|
|
while True:
|
|
try:
|
|
user_input = input("You: ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
print("\nBye!")
|
|
break
|
|
|
|
if not user_input:
|
|
continue
|
|
if user_input.lower() in ("quit", "exit"):
|
|
print("Bye!")
|
|
break
|
|
|
|
response = await run_agent(agent, user_input)
|
|
print(f"Agent: {response}\n")
|
|
|
|
|
|
def main() -> None:
|
|
args = _parse_args()
|
|
load_dotenv()
|
|
config = Config.from_env()
|
|
|
|
if args.model:
|
|
config = Config(
|
|
llm_api_key=config.llm_api_key,
|
|
llm_base_url=config.llm_base_url,
|
|
llm_model=args.model,
|
|
skill_directories=config.skill_directories,
|
|
base_system_prompt=config.base_system_prompt,
|
|
)
|
|
|
|
registry = SkillRegistry()
|
|
|
|
skill_dirs = args.skill_dir or config.skill_directories
|
|
for skill_dir in skill_dirs:
|
|
if skill_dir.is_dir():
|
|
registry.load_from_directory(skill_dir)
|
|
else:
|
|
print(f"Warning: skill directory not found: {skill_dir}", file=sys.stderr)
|
|
|
|
asyncio.run(_run_interactive(config, registry))
|