File size: 3,560 Bytes
1f891e5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
"""
API endpoints for the Legal Assistant service.
"""
import sys
import time
import uuid
import logging
from fastapi import FastAPI, HTTPException, Depends, Request, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import asyncio
from typing import Dict, Optional

# Add AI_core to path so we can import it
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

# Import the process_legal_request function from AI_core
from AI_core.main import process_legal_request

# Import models and config
from .models import QueryRequest, QueryResponse, HealthResponse
from ..config import API_PREFIX, API_VERSION, ALLOWED_ORIGINS

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Create FastAPI app
app = FastAPI(
    title="Legal Assistant API",
    description="API for interacting with the Legal AI Assistant",
    version=API_VERSION,
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Session storage to maintain context
# In production, use a persistent storage solution
sessions: Dict[str, Dict] = {}

@app.get(f"{API_PREFIX}/health", response_model=HealthResponse)
async def health_check():
    """
    Health check endpoint to verify the API is running correctly.
    """
    return HealthResponse(
        status="ok",
        version=API_VERSION,
        components={
            "api": "healthy",
            "agent": "healthy",
            "tools": "healthy"
        }
    )

@app.post(f"{API_PREFIX}/query", response_model=QueryResponse)
async def query(request: QueryRequest):
    """
    Process a legal query and return a response from the Legal AI Assistant.
    """
    # Generate session ID if not provided
    session_id = request.session_id or str(uuid.uuid4())
    
    # Log request
    logger.info(f"Processing query for session {session_id}: {request.query[:50]}...")
    
    # Measure processing time
    start_time = time.time()
    
    try:
        # Process the query using the AI core
        response = await process_legal_request(request.query)
        
        # Calculate processing time
        processing_time = time.time() - start_time
        
        # Return the response
        return QueryResponse(
            response=response,
            session_id=session_id,
            metadata={"tools_used": []},  # In a production version, track tools used
            processing_time=processing_time
        )
    
    except Exception as e:
        logger.error(f"Error processing query: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Error processing your request: {str(e)}"
        )

@app.get(f"{API_PREFIX}/sessions/{{session_id}}")
async def get_session(session_id: str):
    """
    Get information about a specific session.
    """
    if session_id not in sessions:
        raise HTTPException(
            status_code=404,
            detail=f"Session {session_id} not found"
        )
    
    return sessions[session_id]

@app.delete(f"{API_PREFIX}/sessions/{{session_id}}")
async def delete_session(session_id: str):
    """
    Delete a specific session.
    """
    if session_id in sessions:
        del sessions[session_id]
    
    return {"status": "success", "message": f"Session {session_id} deleted"}