File size: 3,351 Bytes
1d21f23
647c8f0
1d21f23
647c8f0
1d21f23
 
647c8f0
1d21f23
647c8f0
1d21f23
647c8f0
1d21f23
 
 
 
 
647c8f0
1d21f23
 
647c8f0
1d21f23
647c8f0
 
1d21f23
647c8f0
 
1d21f23
647c8f0
 
 
 
 
 
1d21f23
647c8f0
1d21f23
 
647c8f0
1d21f23
 
 
 
 
 
647c8f0
1d21f23
 
647c8f0
 
 
1d21f23
 
647c8f0
1d21f23
 
 
647c8f0
 
 
1d21f23
 
647c8f0
 
 
 
1d21f23
 
647c8f0
 
 
1d21f23
 
647c8f0
1d21f23
647c8f0
1d21f23
 
 
 
 
 
 
647c8f0
 
1d21f23
647c8f0
 
 
 
1d21f23
647c8f0
1d21f23
 
 
647c8f0
 
 
1d21f23
 
647c8f0
1d21f23
647c8f0
1d21f23
647c8f0
 
1d21f23
647c8f0
 
 
 
 
1d21f23
647c8f0
1d21f23
 
 
647c8f0
1d21f23
647c8f0
1d21f23
647c8f0
 
 
 
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
import logging
import os
import sys
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, Any, List, Optional
from pathlib import Path
import time

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler("llm_service.log", mode="a"),
    ],
)
logger = logging.getLogger(__name__)

# Import the model
from .model import get_llm_instance

# Initialize model
llm = get_llm_instance()

# Create FastAPI app
app = FastAPI(
    title="LLM Service API",
    description="API for interacting with the local LLM",
    version="1.0.0",
)

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, specify actual origins
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class GenerateRequest(BaseModel):
    prompt: str
    system_prompt: Optional[str] = None
    max_tokens: Optional[int] = 512
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9


class ExpandPromptRequest(BaseModel):
    prompt: str


@app.get("/")
def read_root():
    return {"status": "ok", "message": "LLM Service is running"}


@app.get("/health")
def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "model": "TinyLlama-1.1B-Chat-v1.0"}


@app.post("/generate")
def generate_text(request: GenerateRequest):
    """Generate text from a prompt"""
    try:
        start_time = time.time()
        logger.info(f"Generating text for prompt: {request.prompt[:50]}...")

        result = llm.generate(
            prompt=request.prompt,
            system_prompt=request.system_prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
        )

        elapsed = time.time() - start_time
        logger.info(f"Generation completed in {elapsed:.2f} seconds")

        return {
            "result": result,
            "elapsed_seconds": elapsed,
        }
    except Exception as e:
        logger.error(f"Error during text generation: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/expand-prompt")
def expand_prompt(request: ExpandPromptRequest):
    """Expand a creative prompt with more detail"""
    try:
        start_time = time.time()
        logger.info(f"Expanding prompt: {request.prompt[:50]}...")

        expanded_prompt = llm.expand_creative_prompt(request.prompt)

        elapsed = time.time() - start_time
        logger.info(f"Prompt expansion completed in {elapsed:.2f} seconds")

        return {
            "original_prompt": request.prompt,
            "expanded_prompt": expanded_prompt,
            "elapsed_seconds": elapsed,
        }
    except Exception as e:
        logger.error(f"Error during prompt expansion: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# Run the service when executed directly
if __name__ == "__main__":
    import uvicorn

    port = int(os.environ.get("PORT", 8000))
    host = os.environ.get("HOST", "0.0.0.0")
    logger.info(f"Starting LLM service on {host}:{port}")
    uvicorn.run(app, host=host, port=port)