Agent Workflows

import argparse
import json
import os
from morphik import Morphik
from openai import OpenAI

SYSTEM_INSTRUCTIONS = """You are a helpful assistant with access to a curated knowledge base.

Available metadata fields (customize to your schema):
- doc_id (string)
- doc_type (string)
- title (string)
- category (string)
- status (string)
- effective_date (datetime)
- expires_at (datetime)
- tags (array)
- region (string)

Filter operators:
- $eq, $gt, $gte, $lt, $lte, $and, $contains, $regex

Strategy:
1) Call list_documents to see available material and metadata.
2) Apply filters before retrieval when date, status, region, or tags matter.
3) Use retrieve_chunks with use_colpali=True and output_format="url" to pull text and images.
4) Cite sources with doc_id (or your ID field), filename, and page numbers.
"""

TOOLS = [
    {
        "type": "function",
        "name": "list_documents",
        "description": "List documents in the knowledge base with optional metadata filters.",
        "parameters": {
            "type": "object",
            "properties": {
                "filters": {
                    "type": ["string", "null"],
                    "description": "Metadata filters as a JSON string. Pass null or empty for no filters."
                },
                "limit": {
                    "type": ["integer", "null"],
                    "description": "Maximum number of documents to return (default: 10)"
                }
            },
            "required": ["filters", "limit"],
            "additionalProperties": False
        },
        "strict": True
    },
    {
        "type": "function",
        "name": "retrieve_chunks",
        "description": "Retrieve relevant document chunks for a query using semantic search.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query or user question"
                },
                "filters": {
                    "type": ["string", "null"],
                    "description": "Optional metadata filters as a JSON string"
                },
                "k": {
                    "type": ["integer", "null"],
                    "description": "Number of chunks to retrieve (default: 3)"
                }
            },
            "required": ["query", "filters", "k"],
            "additionalProperties": False
        },
        "strict": True
    }
]

def list_documents(morphik: Morphik, filters=None, limit=None):
    if isinstance(filters, str) and filters.strip():
        filters = json.loads(filters)

    response = morphik.list_documents(filters=filters, limit=limit or 10)
    docs = response.documents if hasattr(response, "documents") else response

    normalized = []
    for doc in docs:
        meta = doc.metadata or {}
        normalized.append({
            "external_id": doc.external_id,
            "filename": doc.filename,
            "doc_id": meta.get("doc_id"),
            "title": meta.get("title"),
            "doc_type": meta.get("doc_type"),
            "category": meta.get("category"),
            "status": meta.get("status"),
            "effective_date": str(meta.get("effective_date")) if meta.get("effective_date") else None,
            "expires_at": str(meta.get("expires_at")) if meta.get("expires_at") else None,
            "tags": meta.get("tags") or [],
        })

    return {"documents": normalized, "count": len(normalized)}


def retrieve_chunks(morphik: Morphik, query, filters=None, k=None):
    if isinstance(filters, str) and filters.strip():
        filters = json.loads(filters)

    chunks = morphik.retrieve_chunks(
        query=query,
        filters=filters,
        k=k or 3,
        use_colpali=True,
        output_format="url"
    )

    results = []
    image_urls = []

    for chunk in chunks:
        meta = chunk.metadata or {}
        is_image_url = isinstance(chunk.content, str) and chunk.content.startswith("http")

        if is_image_url:
            image_urls.append(chunk.content)

        preview = chunk.content
        if isinstance(preview, str) and len(preview) > 300:
            preview = f"{preview[:300]}..."
        if is_image_url:
            preview = f"[Image URL: {chunk.content}]"

        results.append({
            "score": round(chunk.score, 3),
            "content_preview": preview,
            "filename": meta.get("filename") or chunk.filename or "unknown",
            "doc_id": meta.get("doc_id"),
            "doc_type": meta.get("doc_type"),
            "title": meta.get("title"),
            "expires_at": str(meta.get("expires_at")) if meta.get("expires_at") else None,
            "tags": meta.get("tags") or [],
        })

    return {
        "chunks": results,
        "count": len(results),
        "image_urls": image_urls
    }

def run_agent(morphik: Morphik, openai_client: OpenAI, user_query: str, model: str):
    conversation = [{"role": "user", "content": user_query}]

    for _ in range(5):
        response = openai_client.responses.create(
            model=model,
            instructions=SYSTEM_INSTRUCTIONS,
            tools=TOOLS,
            input=conversation
        )

        conversation += response.output
        calls = [item for item in response.output if item.type == "function_call"]

        if not calls:
            print(response.output_text)
            return

        for call in calls:
            args = json.loads(call.arguments)

            if call.name == "list_documents":
                result = list_documents(
                    morphik,
                    filters=args.get("filters"),
                    limit=args.get("limit")
                )
                image_urls = []
            elif call.name == "retrieve_chunks":
                result = retrieve_chunks(
                    morphik,
                    query=args["query"],
                    filters=args.get("filters"),
                    k=args.get("k")
                )
                image_urls = result.get("image_urls", [])
            else:
                result = {"error": f"Unknown function: {call.name}"}
                image_urls = []

            conversation.append({
                "type": "function_call_output",
                "call_id": call.call_id,
                "output": json.dumps(result)
            })

            if image_urls:
                image_content = [{"type": "input_text", "text": "Here are the retrieved pages:"}]
                for url in image_urls:
                    image_content.append({"type": "input_image", "image_url": url})
                conversation.append({"role": "user", "content": image_content})

    print("[Warning] Reached maximum iterations without a final answer")

def main():
    parser = argparse.ArgumentParser(description="Retrieval agent with Morphik tools")
    parser.add_argument("--uri", default=os.getenv("MORPHIK_URI"))
    parser.add_argument("--query", default="What are the latest policy updates published this year?")
    parser.add_argument("--openai-key", default=os.getenv("OPENAI_API_KEY"))
    parser.add_argument("--model", default="gpt-4o-mini")
    args = parser.parse_args()

    if not args.uri:
        raise SystemExit("MORPHIK_URI not set. Use --uri or set MORPHIK_URI.")
    if not args.openai_key:
        raise SystemExit("OPENAI_API_KEY not set. Use --openai-key or set OPENAI_API_KEY.")

    morphik = Morphik(args.uri, timeout=30)
    openai_client = OpenAI(api_key=args.openai_key)
    run_agent(morphik, openai_client, args.query, args.model)


if __name__ == "__main__":
    main()

export MORPHIK_URI="morphik://your-token@your-host"
export OPENAI_API_KEY="your-openai-key"
python retrieval_agent.py --query "Where can I find the most recent safety policies?"