kasinathansj/Categorization

1import logging2import time3import json4import re5from fastapi import FastAPI, HTTPException6from pydantic import BaseModel7import uvicorn8from llama_cpp import Llama9from severity_check import get_highest_severity_word, SEVERITY_ORDER10from producer import send_to_repetitive_queue11import httpx12import os13import psutil14from typing import Dict, Any15from huggingface_hub import hf_hub_download16 17 18# Setup logging19logging.basicConfig(20    level=logging.INFO,21    format="%(asctime)s - %(levelname)s - %(message)s",22    handlers=[23        logging.StreamHandler(),  # Logs to console24        logging.FileHandler("server.log", mode="a")  # Logs to file25    ]26)27logging.getLogger("httpx").setLevel(logging.WARNING)28 29app = FastAPI()30 31# Modify input model to include organizationId instead of organizationName32class TextInput(BaseModel):33    text: str34    organizationId: str35    petitionDetails: Dict[str, Any]  # Dictionary to store petition details36 37# Fetch departments based on organizationId using the provided URL format.38async def fetch_departments(org_id: str):39    # Example: http://localhost:4000/organizations/{org_id}/departments40    const_base_url = os.getenv("PETITION_SERVICE_BACKEND_URL", "http://localhost:4000")41    url = f"{const_base_url}/organizations/{org_id}/departments"42    print(url)43    try:44        async with httpx.AsyncClient() as client:45            response = await client.get(url)46            response.raise_for_status()47            data = response.json()["departments"]48            return [49                {50                    "department": dept["name"],51                    "description": dept["description"],52                    "id": dept["id"]53                }54                for dept in data55            ]56    except Exception as e:57        logging.error(f"Failed to fetch departments: {e}")58        raise HTTPException(status_code=500, detail="Failed to fetch departments")59 60@app.get("/cpu-usage")61def get_cpu_usage():62    """Returns the current CPU usage percentage."""63    try:64        cpu_usage = psutil.cpu_percent(interval=1)  # 1-second interval65        return {"cpu_usage": cpu_usage}66    except Exception as e:67        logging.error(f"Failed to retrieve CPU usage: {e}")68        raise HTTPException(status_code=500, detail="Failed to retrieve CPU usage")69 70@app.get("/")71def home():72    logging.info("Home endpoint accessed")73    return {"message": "Categorization server is running!"}74 75 76# Define model details77# REPO_ID = "TheBloke/Mistral-7B-Instruct-v0.2-GGUF"78# FILENAME = "mistral-7b-instruct-v0.2.Q4_0.gguf"79 80REPO_ID = "TheBloke/Mistral-7B-Instruct-v0.1-GGUF"81FILENAME = "mistral-7b-instruct-v0.1.Q4_K_M.gguf" 82 83def load_model():84    try:85        logging.info("Downloading & Loading model...")86 87        # Automatically download from Hugging Face Hub88        model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)89 90        # Load model91        model = Llama(92            model_path=model_path,93            n_ctx=8192,  # Increase if the model supports it94            n_threads=4,  # Adjust for CPU95            n_batch=512,  # Tune for performance96        )97 98        logging.info("✅ Model loaded successfully!")99        return model100 101    except Exception as e:102        logging.error(f"❌ Error loading model: {e}")103        raise RuntimeError("Failed to load model")104 105llm = load_model()106 107# Generate classification prompt108def generate_prompt(text, departments):109    converted_department = [110        {111            "department": dept["department"],112            "description": dept["description"],113        }114        for dept in departments115    ]116    department_definitions = json.dumps(converted_department, indent=4)117    118    return f"""119    You are an AI that strictly returns a JSON object categorizing text into the most relevant department.120 121    **Departments and Definitions:**122    {department_definitions}123 124    **Follow these rules:**125    - Always return a **single** department that best fits.126    - Assess the **severity** as one of the following: "High", "Medium", or "Low".127    - Do not add extra text before or after the JSON.128    - Stick to the format exactly as shown.129 130    **Example Output:**131    132    {{133        "department": "Example Department",134        "reason": "The text pertains to example-related matters.",135        "severity": "High"136    }}137 138    **Classify the following text with reason and determine its severity:**139 140    "{text}"141    """.strip()142 143# Extract JSON response from model output144def extract_json(response_text):145    json_match = re.search(r"\{.*\}", response_text, re.DOTALL)146    if json_match:147        response_text = json_match.group(0)148    try:149        return json.loads(response_text)150    except json.JSONDecodeError:151        logging.error(f"Failed to parse JSON: {response_text}")152        return {"department": "Unknown", "reason": "Response not in expected format.", "severity": "Unknown"}153 154# Function to classify a single text155def classify_petition(text, departments):156    logging.info(f"Classifying text: {text}")157 158    local_severity = get_highest_severity_word(text)159    prompt = generate_prompt(text, departments)160 161    try:162        start_time = time.time()163        output = llm(prompt, max_tokens=256, temperature=0.2)164        execution_time = round(time.time() - start_time, 2)165 166        response_text = output["choices"][0]["text"].strip()167        response_json = extract_json(response_text)168 169        department = response_json.get("department", "Unknown")170        reason = response_json.get("reason", "No reason provided.")171        llm_severity = response_json.get("severity", "Unknown")172        logging.info(f"Classification result: {department} (Time: {execution_time}s)")173        if local_severity is None:174            final_severity = llm_severity  # Default to LLM severity if no local severity is found175        elif llm_severity is None:176            final_severity = local_severity  # Default to local severity if LLM is missing177        else:178            # Both severities exist, compare them179            final_severity = local_severity if SEVERITY_ORDER.get(local_severity, 0) >= SEVERITY_ORDER.get(llm_severity, 0) else llm_severity180        return {181            "text": text,182            "department": department,183            "severity": final_severity,184            "reason": reason,185            "time_taken": execution_time186        }187    except Exception as e:188        logging.error(f"Error during classification: {e}")189        return {"error": "Internal classification error"}190 191@app.post("/classify")192async def classify(text_input: TextInput):193    logging.info(f"Received classification request: {text_input.text}")194 195    try:196        # Now fetch departments using organizationId197        departments = await fetch_departments(text_input.organizationId)198        result = classify_petition(text_input.text, departments)199        if result["department"]:200            matching_department = next(201                (dept for dept in departments if dept["department"] == result["department"]), 202                None203            )204            if matching_department:205                result["departmentId"] = matching_department["id"]206        result["organizationId"] = text_input.organizationId207        result["petitionDetails"] = text_input.petitionDetails208        logging.info(f"Response: {result}")209        send_to_repetitive_queue(result)  # Send to Kafka queue210        211        return result212    except Exception as e:213        logging.error(f"Classification request failed: {e}")214        raise HTTPException(status_code=500, detail="Internal server error")215 216if __name__ == "__main__":217    logging.info("Starting FastAPI server...")218    uvicorn.run(app, host="0.0.0.0", port=8000)219