New hire onboarding agent with citations

import os
import re
import time

from seekrai import SeekrFlow
from seekrai.types import CreateAgentRequest, FileSearch, FileSearchEnv

API_KEY = os.environ["SEEKR_API_KEY"]
API_URL = "https://flow.seekr.com/v1/"
MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
EMBEDDING_MODEL = "intfloat/e5-mistral-7b-instruct"

FILE_PATHS = [
    "company-guidebook.pdf",
    "company-holidays.pdf",
    "company-payroll-schedule.pdf",
]

client = SeekrFlow(api_key=API_KEY, base_url=API_URL)

vector_db = client.vector_database.create(
    name="BennyBotDB",
    model=EMBEDDING_MODEL,
    description="Employee guidelines and benefits for onboarding QA.",
)
database_id = vector_db.id
print(f"Created vector database: {vector_db.name} ({database_id})")

file_ids = []
for file_path in FILE_PATHS:
    print(f"Uploading {file_path}...")
    upload_response = client.files.upload(file_path, purpose="alignment")
    file_ids.append(upload_response.id)
    print(f"  Uploaded: {upload_response.id}")

ingestion_job = client.vector_database.create_ingestion_job(
    database_id=database_id,
    files=file_ids,
    method="accuracy-optimized",
    chunking_method="markdown",
    token_count=512,
    overlap_tokens=50,
)
job_id = ingestion_job.id
print(f"Created ingestion job: {job_id}")

deadline = time.time() + 600
while time.time() < deadline:
    job = client.vector_database.retrieve_ingestion_job(database_id, job_id)
    print(f"Ingestion status: {job.status}")
    if job.status == "completed":
        print("Vector database ready.")
        break
    if job.status == "failed":
        raise RuntimeError(f"Ingestion failed: {job.error_message}")
    time.sleep(5)
else:
    raise TimeoutError("Ingestion did not complete in time.")

instructions = """You are an expert onboarding assistant that answers only from document search results.

For each question:
1. Search the documents with the file_search tool.
2. Use information from multiple sources when available.
3. When you find relevant information, rate your confidence from 1 (guess) to 10 (certain).
4. After your answer, on a new line starting with "Confidence: [X/10]", briefly explain the rating.
5. Cite your sources, including specific document names.

If the answer is not in the search results, say you could not find it and do not give a confidence score."""

agent = client.agents.create(
    CreateAgentRequest(
        name="BennyBot",
        instructions=instructions,
        model_id=MODEL_ID,
        tools=[
            FileSearch(
                tool_env=FileSearchEnv(
                    file_search_index=database_id,
                    document_tool_desc="Search onboarding documents to answer employee questions accurately.",
                    top_k=6,
                    score_threshold=0.5,
                )
            )
        ],
    )
)
agent_id = agent.id
print(f"Agent created: {agent_id}")

deadline = time.time() + 300
while time.time() < deadline:
    status = client.agents.retrieve(agent_id=agent_id).status.value
    print(f"Agent status: {status}")
    if status == "Active":
        print("Agent is active.")
        break
    if status == "Failed":
        raise RuntimeError("Agent failed to deploy.")
    time.sleep(10)
else:
    raise TimeoutError("Agent did not become active in time.")

def parse_response(response: str) -> dict:
    """Split an agent reply into answer, confidence, explanation, and sources."""
    if not response:
        return {"answer": "No response received.", "confidence": "0/10", "explanation": "", "sources": []}

    answer, confidence, explanation, sources = response, "Not provided", "", []

    if "Confidence:" in response:
        answer, _, confidence_part = response.partition("Confidence:")
        answer = answer.strip()
        confidence_part = confidence_part.strip()
        rating_match = re.search(r"(\d+)(/10)?", confidence_part)
        if rating_match:
            confidence = f"{rating_match.group(1)}/10"
            explanation = re.sub(r"^\d+(/10)?", "", confidence_part).strip()

    for pattern in (r"Source[s]?:(.+?)(?=\n\n|\Z)", r"According to (.+?)(?=\n\n|\Z)"):
        match = re.search(pattern, response, re.IGNORECASE | re.DOTALL)
        if match:
            sources += [s.strip() for s in re.split(r",|\n", match.group(1)) if s.strip()]

    return {"answer": answer, "confidence": confidence, "explanation": explanation, "sources": sources}


thread = client.agents.threads.create()
client.agents.threads.create_message(
    thread_id=thread.id, role="user", content="How many paid holidays do employees get?"
)
run_response = client.agents.runs.run(agent_id=agent_id, thread_id=thread.id, stream=False)

deadline = time.time() + 300
while time.time() < deadline:
    run = client.agents.runs.retrieve(run_response.run_id, thread.id)
    if run.status.value == "completed":
        break
    if run.status.value in ("failed", "canceled"):
        raise RuntimeError(f"Run ended with status: {run.status.value}")
    time.sleep(2)
else:
    raise TimeoutError("Run did not complete in time.")

messages = client.agents.threads.list_messages(thread.id, limit=10, order="desc")
reply = next((m.content for m in messages if m.role == "assistant"), None)

result = parse_response(reply if isinstance(reply, str) else str(reply))
print("\n" + "=" * 50)
print(f"Answer: {result['answer']}\n")
print(f"Confidence: {result['confidence']}")
if result["explanation"]:
    print(f"Explanation: {result['explanation']}")
if result["sources"]:
    print("\nSources:")
    for i, source in enumerate(result["sources"], 1):
        print(f"{i}. {source}")