Real Estate Agent with LangGraph (using CLI)

# helper packages
import os
import time
import json
import logging
from typing import Optional, Type, Annotated, TypedDict, Union, List
from pydantic import Field, BaseModel, Json
import asyncio
from pathlib import Path
from dotenv import load_dotenv

# LangGraph and LangChain imports
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_core.tools import StructuredTool

# TensorLake imports
from tensorlake.documentai import DocumentAI, ParsingOptions
from tensorlake.documentai.parse import (ExtractionOptions)

load_dotenv()

# Load environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
TENSORLAKE_API_KEY = os.getenv("TENSORLAKE_API_KEY")

def detect_signatures_in_document(path: str) -> str:
    """
    Uploads a document to TensorLake, triggers parsing, including signature detection, and returns the parsed result.
    """
    if not Path(path).exists():
        return f"File not found: {path}"

    if not TENSORLAKE_API_KEY:
        return "Error: TENSORLAKE_API_KEY environment variable is not set"

    # Initialize DocumentAI client
    doc_ai = DocumentAI(api_key=TENSORLAKE_API_KEY)
    
    # Upload document to TensorLake
    file_id = doc_ai.upload(path=path)

    parsing_options = ParsingOptions(
        detect_signature=True,  # this needs to be True
        extraction_options=ExtractionOptions(skip_ocr=True),
    )
    # Start parsing job
    job_id = doc_ai.parse(file_id, options=parsing_options)

    # Poll for completion with configurable timeout
    start_time = time.time()
    max_wait_time = 300  # or set as a constant or parameter (here, we will wait for 5 min max)

    while time.time() - start_time < max_wait_time:
        # Signature detection result after parsing the document
        result = doc_ai.get_job(job_id)  # this may take 2-3 minutes

        if result.status in ["pending", "processing"]:
            time.sleep(5)  # Wait 5 seconds before checking again
        elif result.status == "successful":

            # Optional: save the parsed result in a file for referring to it later
            with open(f"parsed_result_{file_id}.json", "w") as f:
                json.dump(result.model_dump(), f, indent=2)

            # Return parsed result
            return str(result)
        else:
            return f"Document parsing failed with status: {result.status}"

    # Timeout reached
    return f"Document processing timeout after {max_wait_time} seconds. Job ID: {job_id}"

async def detect_signatures_in_document_async(path: str) -> str:
    """Asynchronous version of detecting signatures from document."""
    return await asyncio.to_thread(detect_signatures_in_document, path)

# Create the LangChain tool using StructuredTool
signature_detection_tool = StructuredTool.from_function(
    func=detect_signatures_in_document,
    coroutine=detect_signatures_in_document_async,
    name="SignatureDetectionTool",
    description="Detect signatures from any document (PDF, Markdown, Docx, etc.)",
    return_direct=False,
    handle_tool_error="Document parsing failed. Please verify the file path and your Tensorlake API key."
)

def build_document_analysis_prompt(parsed_result: str, questions: Union[str, List[str]]) -> str:
    # Normalize single question to list
    if isinstance(questions, str):
        questions = [questions]

    question_block = "\n".join(f"{i + 1}. {q}" for i, q in enumerate(questions))

    system_prompt = f"""You are a helpful assistant that answers questions about documents with signature detection data.

Your responsibilities:
1. Answer questions based on that loaded data
2. Help users understand the signature analysis results

You can answer questions like:
- How many signatures were found?
- Which pages contain signatures?
- Who signed the document?
- What does the content say around signatures?
- What type of document is this?
- Who are the parties involved?
- What is the date of the signature?
- Did each party sign the document?
- Are there any missing signatures on any pages?
- Which property is missing signatures?
- Who is the agent for the properties missing signatures?

I've processed a document and got this result:
{parsed_result}

Please analyze the above parsed output and answer the following:
{question_block}
"""
    return system_prompt

# Define the agent state
class AgentState(TypedDict):
    messages: Annotated[list, add_messages]


# Agent node - decides whether to use tools
async def agent_node(state: AgentState):
    model = ChatOpenAI(
        model="gpt-4o",
        temperature=0.1
    ).bind_tools([signature_detection_tool])

    response = await model.ainvoke(state["messages"])
    return {"messages": [response]}


# Conditional Logic for Tool Use
def should_continue(state: AgentState):
    last_message = state["messages"][-1]
    if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        return "tools"
    return END

# LangGraph Workflow Setup
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
workflow.add_node("tools", ToolNode([signature_detection_tool]))
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", should_continue, {"tools": "tools", END: END})
workflow.add_edge("tools", "agent")
app = workflow.compile()

# Document + Agent Pipeline
async def analyze_signatures_agents(
        path: str,
        questions: List[str]
) -> str:
    """Invoke the tool with parsing options, then use agent for analysis."""

    print("🔍 Processing document with signature detection...")

    # Pass parsing options and run the tool
    parsed_output = await signature_detection_tool.ainvoke({
        "path": path
    })

    # Build prompt
    prompt = build_document_analysis_prompt(parsed_output, questions)

    # Run agent on prompt
    final_state = await app.ainvoke({
        "messages": [HumanMessage(content=prompt)]
    })

    return final_state["messages"][-1].content

async def example_signature_detection_real_estate():

    # change to your own file path
    path = "path/to/your/document.pdf"

    analysis_questions = [
        "How many signatures were detected in this document and who are the parties involved?",
        "What contextual information can you extract about any signatures?",
        "Are there any missing signatures on any pages?"
    ]

    result = await analyze_signatures_agents(
        path=path,
        questions=analysis_questions
    )

    print("Analysis Result:\n\n", result)


if __name__ == "__main__":
    # run the example
    asyncio.run(example_signature_detection_real_estate())