Building a Legal Practice Memory System

from papr_memory import Papr
import os

client = Papr(x_api_key=os.environ.get("PAPR_MEMORY_API_KEY"))

# Define legal contract schema
legal_schema = client.schemas.create(
    name="Legal Contract Schema",
    description="Schema for managing legal contracts, parties, obligations, and case notes",
    version="1.0.0",
    node_types={
        "Contract": {
            "name": "Contract",
            "label": "Contract",
            "description": "Legal contract or agreement document",
            "properties": {
                "title": {
                    "type": "string",
                    "required": True,
                    "description": "Full contract title as it appears in the document"
                },
                "contract_type": {
                    "type": "string",
                    "required": True,
                    "description": "Type of legal contract",
                    "enum_values": ["service", "employment", "nda", "partnership", "license"]
                },
                "value": {
                    "type": "float",
                    "required": False,
                    "description": "Total contract value in USD"
                },
                "effective_date": {
                    "type": "datetime",
                    "required": False,
                    "description": "Date when contract becomes effective"
                },
                "expiration_date": {
                    "type": "datetime",
                    "required": False,
                    "description": "Date when contract expires"
                },
                "status": {
                    "type": "string",
                    "required": False,
                    "description": "Current status of the contract",
                    "enum_values": ["draft", "negotiating", "active", "expired", "terminated"],
                    "default": "draft"
                },
                "client_id": {
                    "type": "string",
                    "required": False,
                    "description": "Client identifier for this contract"
                }
            },
            "required_properties": ["title", "contract_type"],
            "unique_identifiers": ["title"],
            "color": "#e74c3c"
        },
        "Party": {
            "name": "Party",
            "label": "Party",
            "description": "Legal party to a contract",
            "properties": {
                "name": {
                    "type": "string",
                    "required": True,
                    "description": "Full legal name of the party"
                },
                "party_type": {
                    "type": "string",
                    "required": False,
                    "description": "Type of party",
                    "enum_values": ["individual", "company", "government", "nonprofit"]
                },
                "role": {
                    "type": "string",
                    "required": False,
                    "description": "Role in the contract",
                    "enum_values": ["client", "counterparty", "vendor", "partner"]
                }
            },
            "required_properties": ["name"],
            "unique_identifiers": ["name"],
            "color": "#3498db"
        },
        "Obligation": {
            "name": "Obligation",
            "label": "Obligation",
            "description": "Contractual obligation or deadline",
            "properties": {
                "description": {
                    "type": "string",
                    "required": True,
                    "description": "Description of the obligation"
                },
                "deadline": {
                    "type": "datetime",
                    "required": False,
                    "description": "Deadline for fulfilling the obligation"
                },
                "status": {
                    "type": "string",
                    "required": False,
                    "description": "Current status",
                    "enum_values": ["pending", "completed", "overdue"],
                    "default": "pending"
                },
                "responsible_party": {
                    "type": "string",
                    "required": False,
                    "description": "Party responsible for this obligation"
                }
            },
            "required_properties": ["description"],
            "unique_identifiers": []
        },
        "CaseNote": {
            "name": "CaseNote",
            "label": "Case Note",
            "description": "Lawyer's notes about a case or contract",
            "properties": {
                "note": {
                    "type": "string",
                    "required": True,
                    "description": "Note content"
                },
                "note_type": {
                    "type": "string",
                    "required": False,
                    "description": "Type of note",
                    "enum_values": ["negotiation", "client_preference", "risk", "precedent"]
                },
                "created_date": {
                    "type": "datetime",
                    "required": False,
                    "description": "Date note was created"
                }
            },
            "required_properties": ["note"],
            "unique_identifiers": []
        }
    },
    relationship_types={
        "PARTY_TO": {
            "name": "PARTY_TO",
            "label": "Party To",
            "description": "Party is signatory to the contract",
            "allowed_source_types": ["Party"],
            "allowed_target_types": ["Contract"]
        },
        "HAS_OBLIGATION": {
            "name": "HAS_OBLIGATION",
            "label": "Has Obligation",
            "description": "Contract includes an obligation",
            "allowed_source_types": ["Contract"],
            "allowed_target_types": ["Obligation"]
        },
        "NOTE_ABOUT": {
            "name": "NOTE_ABOUT",
            "label": "Note About",
            "description": "Note is about a contract",
            "allowed_source_types": ["CaseNote"],
            "allowed_target_types": ["Contract"]
        }
    }
)

schema_id = legal_schema.data.id
print(f"Legal schema created: {schema_id}")

# Upload a service agreement
doc_response = client.document.upload(
    file=open("acme_service_agreement.pdf", "rb"),
    schema_id=schema_id,
    simple_schema_mode=True,
    hierarchical_enabled=True,
    user_id="lawyer_sarah",  # Associate with Sarah's workspace
    property_overrides=[
        {
            "nodeLabel": "Contract",
            "set": {
                "status": "negotiating",
                "client_id": "acme_corp"
            }
        },
        {
            "nodeLabel": "Party",
            "match": {"name": "Acme Corp"},
            "set": {
                "id": "party_acme",
                "role": "client"
            }
        }
    ]
)

upload_id = doc_response.document_status.upload_id
print(f"Document uploaded: {upload_id}")

# Monitor processing
import time
while True:
    status = client.document.get_status(upload_id)
    print(f"Progress: {status.progress * 100}%")
    
    if status.status_type == "completed":
        print("✓ Contract processed successfully!")
        break
    elif status.status_type == "failed":
        print(f"✗ Processing failed: {status.error}")
        break
    
    time.sleep(5)

# Add negotiation note
negotiation_note = client.memory.add(
    content="Acme Corp negotiation: Client strongly prefers termination for convenience clause. Previous contracts show they value flexibility over cost savings. Suggested 30-day notice period was well received.",
    user_id="lawyer_sarah",
    metadata={
        "client_id": "acme_corp",
        "contract_type": "service",
        "topics": ["negotiation", "termination_clause", "client_preference"]
    },
    graph_generation={
        "mode": "auto",
        "auto": {
            "schema_id": schema_id,
            "simple_schema_mode": True
        }
    }
)

# Add risk assessment note
risk_note = client.memory.add(
    content="Risk assessment for Acme service agreement: Payment terms of Net-60 are longer than our standard Net-30. Recommend adding late payment penalties of 1.5% monthly interest. Client has good payment history based on previous contracts.",
    user_id="lawyer_sarah",
    metadata={
        "client_id": "acme_corp",
        "contract_type": "service",
        "topics": ["risk", "payment_terms", "client_history"]
    },
    graph_generation={
        "mode": "auto",
        "auto": {
            "schema_id": schema_id,
            "simple_schema_mode": True
        }
    }
)

# Add precedent note
precedent_note = client.memory.add(
    content="Similar termination clause language used successfully in TechCo agreement (2023). Client accepted 30-day notice with mutual termination rights. Language can be reused as precedent.",
    user_id="lawyer_sarah",
    metadata={
        "topics": ["precedent", "termination_clause", "contract_template"]
    },
    graph_generation={
        "mode": "auto",
        "auto": {
            "schema_id": schema_id,
            "simple_schema_mode": True
        }
    }
)

print("✓ Notes added to memory system")

# Find contracts expiring soon
expiring_response = client.memory.search(
    query="What contracts are expiring in the next 90 days?",
    user_id="lawyer_sarah",
    enable_agentic_graph=True,
    max_memories=20,
    max_nodes=15
)

print(f"Found {len(expiring_response.data.memories)} contracts expiring soon")
for node in expiring_response.data.nodes:
    if node.label == "Contract":
        print(f"\n{node.properties.get('title')}")
        print(f"  Expires: {node.properties.get('expiration_date')}")
        print(f"  Client: {node.properties.get('client_id')}")

# Find precedents for termination clauses
precedent_response = client.memory.search(
    query="Show me precedents for termination clauses with 30-day notice",
    user_id="lawyer_sarah",
    enable_agentic_graph=True,
    max_memories=10
)

print(f"\n\nFound {len(precedent_response.data.memories)} relevant precedents")
for memory in precedent_response.data.memories:
    print(f"\n{memory.content[:200]}...")

# Find client preferences
preference_response = client.memory.search(
    query="What are Acme Corp's contract preferences and negotiation history?",
    user_id="lawyer_sarah",
    metadata_filter={"client_id": "acme_corp"},
    enable_agentic_graph=True,
    max_memories=10
)

print(f"\n\nFound {len(preference_response.data.memories)} client preferences")
for memory in preference_response.data.memories:
    print(f"- {memory.content[:150]}...")

# Find obligations due this month
obligations_response = client.memory.search(
    query="What contract obligations are due this month?",
    user_id="lawyer_sarah",
    enable_agentic_graph=True,
    max_memories=20,
    max_nodes=20
)

print(f"\n\nFound {len(obligations_response.data.nodes)} obligations")
for node in obligations_response.data.nodes:
    if node.label == "Obligation":
        print(f"- {node.properties.get('description')}")
        print(f"  Deadline: {node.properties.get('deadline')}")
        print(f"  Status: {node.properties.get('status')}")

# Get portfolio summary
portfolio_summary = client.graphql.query(
    query="""
    query PortfolioSummary($userId: String!) {
      contracts(where: { user_id: $userId }) {
        total: aggregate { count }
        by_type {
          contract_type
          count: aggregate { count }
          total_value: aggregate { sum(field: "value") }
        }
        by_status {
          status
          count: aggregate { count }
        }
        by_client {
          client_id
          count: aggregate { count }
          total_value: aggregate { sum(field: "value") }
        }
      }
    }
    """,
    variables={"userId": "lawyer_sarah"}
)

print("\n=== Portfolio Summary ===")
summary = portfolio_summary.data['contracts']
print(f"Total contracts: {summary['total']['count']}")

print("\nBy Type:")
for type_data in summary['by_type']:
    print(f"  {type_data['contract_type']}: {type_data['count']} contracts")
    print(f"    Total value: ${type_data['total_value']:,.2f}")

print("\nBy Status:")
for status_data in summary['by_status']:
    print(f"  {status_data['status']}: {status_data['count']} contracts")

# Find upcoming deadlines
upcoming_obligations = client.graphql.query(
    query="""
    query UpcomingDeadlines($userId: String!, $startDate: DateTime!, $endDate: DateTime!) {
      obligations(
        where: {
          contract: { user_id: $userId }
          deadline: { _gte: $startDate, _lte: $endDate }
          status: "pending"
        }
        order_by: { deadline: asc }
      ) {
        description
        deadline
        responsible_party
        contract {
          title
          client_id
        }
      }
    }
    """,
    variables={
        "userId": "lawyer_sarah",
        "startDate": "2024-01-01T00:00:00Z",
        "endDate": "2024-03-31T23:59:59Z"
    }
)

print("\n=== Upcoming Obligations (Next 90 Days) ===")
for obligation in upcoming_obligations.data['obligations']:
    print(f"\n{obligation['description']}")
    print(f"  Deadline: {obligation['deadline']}")
    print(f"  Contract: {obligation['contract']['title']}")
    print(f"  Client: {obligation['contract']['client_id']}")

# Analyze client relationships
client_analysis = client.graphql.query(
    query="""
    query ClientAnalysis($userId: String!) {
      parties(
        where: { 
          role: "client"
          contracts: { user_id: $userId }
        }
      ) {
        name
        contract_count: contracts_aggregate { count }
        total_value: contracts_aggregate { 
          sum { value } 
        }
        active_contracts: contracts(
          where: { status: "active" }
        ) {
          title
          contract_type
          expiration_date
        }
        notes: case_notes {
          note
          note_type
        }
      }
    }
    """,
    variables={"userId": "lawyer_sarah"}
)

print("\n=== Client Relationship Analysis ===")
for client in client_analysis.data['parties']:
    print(f"\n{client['name']}")
    print(f"  Total contracts: {client['contract_count']}")
    print(f"  Total value: ${client['total_value']['sum']:,.2f}")
    print(f"  Active contracts: {len(client['active_contracts'])}")
    if client['notes']:
        print(f"  Key notes:")
        for note in client['notes'][:3]:
            print(f"    - [{note['note_type']}] {note['note'][:100]}...")

class LegalPracticeMemory:
    def __init__(self, api_key: str, lawyer_id: str, schema_id: str):
        self.client = Papr(x_api_key=api_key)
        self.lawyer_id = lawyer_id
        self.schema_id = schema_id
    
    def upload_contract(self, file_path: str, client_id: str, contract_type: str):
        """Upload and process a contract"""
        with open(file_path, "rb") as f:
            response = self.client.document.upload(
                file=f,
                schema_id=self.schema_id,
                simple_schema_mode=True,
                hierarchical_enabled=True,
                user_id=self.lawyer_id,
                property_overrides=[
                    {
                        "nodeLabel": "Contract",
                        "set": {"client_id": client_id, "contract_type": contract_type}
                    }
                ]
            )
        return response.document_status.upload_id
    
    def add_case_note(self, content: str, client_id: str, note_type: str, topics: list):
        """Add a case note"""
        return self.client.memory.add(
            content=content,
            user_id=self.lawyer_id,
            metadata={
                "client_id": client_id,
                "note_type": note_type,
                "topics": topics
            },
            graph_generation={
                "mode": "auto",
                "auto": {
                    "schema_id": self.schema_id,
                    "simple_schema_mode": True
                }
            }
        )
    
    def find_precedents(self, query: str):
        """Find relevant precedents"""
        response = self.client.memory.search(
            query=query,
            user_id=self.lawyer_id,
            metadata_filter={"note_type": "precedent"},
            enable_agentic_graph=True,
            max_memories=10
        )
        return response.data.memories
    
    def get_client_history(self, client_id: str):
        """Get complete client history"""
        response = self.client.memory.search(
            query=f"Complete history and notes for {client_id}",
            user_id=self.lawyer_id,
            metadata_filter={"client_id": client_id},
            enable_agentic_graph=True,
            max_memories=50,
            max_nodes=30
        )
        return {
            "memories": response.data.memories,
            "entities": response.data.nodes
        }
    
    def get_expiring_contracts(self, days: int = 90):
        """Get contracts expiring within specified days"""
        response = self.client.graphql.query(
            query="""
            query ExpiringContracts($userId: String!, $cutoffDate: DateTime!) {
              contracts(
                where: {
                  user_id: $userId
                  expiration_date: { _lte: $cutoffDate }
                  status: "active"
                }
                order_by: { expiration_date: asc }
              ) {
                title
                client_id
                expiration_date
                value
                parties {
                  name
                  role
                }
              }
            }
            """,
            variables={
                "userId": self.lawyer_id,
                "cutoffDate": f"2024-{(datetime.now() + timedelta(days=days)).strftime('%m-%d')}T00:00:00Z"
            }
        )
        return response.data['contracts']

# Usage
practice = LegalPracticeMemory(
    api_key=os.environ.get("PAPR_MEMORY_API_KEY"),
    lawyer_id="lawyer_sarah",
    schema_id=schema_id
)

# Upload new contract
upload_id = practice.upload_contract(
    "new_client_agreement.pdf",
    client_id="newcorp",
    contract_type="service"
)

# Add negotiation notes
practice.add_case_note(
    content="NewCorp prefers aggressive SLAs and is willing to pay premium for 99.9% uptime guarantee",
    client_id="newcorp",
    note_type="negotiation",
    topics=["sla", "pricing", "client_preference"]
)

# Find similar precedents
precedents = practice.find_precedents("SLA clauses with 99.9% uptime guarantee")
print(f"Found {len(precedents)} similar precedents")

# Check expiring contracts
expiring = practice.get_expiring_contracts(days=60)
print(f"\n{len(expiring)} contracts expiring in next 60 days")
for contract in expiring:
    print(f"  - {contract['title']} (Client: {contract['client_id']})")
    print(f"    Expires: {contract['expiration_date']}")