main.py

from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import json
import re
import os
from llm.common import LlmParams, LlmPredictParams
from llm.deepinfra_api import DeepInfraApi
from llm import prompts
from prompts import gettable
from dotenv import load_dotenv
import uvicorn

# Загрузка переменных окружения из файла .env
load_dotenv()

LLM_API_URL = os.getenv("LLM_API_URL", "https://api.deepinfra.com")
LLM_API_KEY = os.getenv("DEEPINFRA_API_KEY", "")
LLM_NAME = os.getenv("LLM_NAME", "meta-llama/Llama-3.3-70B-Instruct-Turbo")

default_llm_params = LlmParams(
    url=LLM_API_URL,
    api_key=LLM_API_KEY,
    model=LLM_NAME,
    predict_params=LlmPredictParams(
        temperature=0.15, top_p=0.95, min_p=0.05, seed=42,
        repetition_penalty=1.2, presence_penalty=1.1, max_tokens=6000
    )
)
llm_api = DeepInfraApi(default_llm_params)

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"]
)

class TextRequest(BaseModel):
    text: str
    projects: list[str] = []

@app.post("/extracttable")
async def extracttable_route(request: TextRequest):
    prompt = gettable.USER_PROMPT.format(query=request.text)
    response = await llm_api.predict(prompt[:150000], system_prompt=gettable.SYSTEM_PROMPT)
    
    result = {"response": None, "error": None, "raw": response}
    
    if "JSON: " not in response:
        result["error"] = "Строка не содержит 'JSON: '"
        return result
    
    prefix, json_str = response.split("JSON: ", 1)
    json_str = json_str.strip()
    
    if not json_str:
        result["error"] = "После 'JSON: ' отсутствует JSON"
        return result
    
    try:
        result["response"] = json.loads(json_str)
        result["raw"] = prefix.strip()
    except json.JSONDecodeError as e:
        result["error"] = f"Ошибка декодирования JSON: {e}"
        
    return result

@app.get("/health")
def health():
    return {"status": "ok"}

async def generate_response(prompt):
    return await llm_api.predict(prompt[:150000])

@app.post("/getsummary")
async def getsummary_route(request: TextRequest):
    return {"result": await generate_response(prompts.GET_SUMMARY.format(text=request.text))}

@app.post("/cleantext")
async def cleantext_route(request: TextRequest):
    return {"result": await generate_response(prompts.CLEAN_TEXT.format(text=request.text))}

@app.post("/getfollowup")
async def getfollowup_route(request: TextRequest):
    return {"result": await generate_response(prompts.GET_FOLLOWUP.format(text=request.text))}

@app.post("/getagenda")
async def getagenda_route(request: TextRequest):
    return {"result": await generate_response(prompts.GET_AGENDA.format(text=request.text))}

@app.post("/gethighlights")
async def gethighlights_route(request: TextRequest):
    return {"result": await generate_response(prompts.GET_HIGHLIGHTS.format(text=request.text))}

@app.post("/getprojectinfo")
async def getprojectinfo_route(request: TextRequest):
    prompts_list = [
        prompts.GET_PROJECT_INFO_NAMES.format(text=request.text),
        prompts.GET_PROJECT_INFO_AGENDA.format(text=request.text)
    ]
    main_info = "\n\n".join([await generate_response(p) for p in prompts_list])
    return {"result": main_info.strip().replace("Конец ответа", "").replace('</s>', '')}

@app.post("/getprojectlist")
async def getprojectlist_route(request: TextRequest):
    list_of_projects = await generate_response(prompts.GET_PROJECT_LIST.format(text=request.text))
    projects = [f"Проект {x}" for x in list_of_projects.split("Проект ") if x][1:]
    projects = [p.replace("проект ", "").strip() for p in projects]
    real_projects = await generate_response(prompts.GET_PROJECT_LIST_CHECK_PROJECT.format(text=request.text, projects=projects))
    real_projects_list = re.findall(r'Да:\s*(.*?)\s*(?:\n\n|$)', real_projects)
    return {"result": real_projects_list}

@app.post("/getprojectdetails")
async def getprojectdetails_route(request: TextRequest):
    if not request.projects:
        raise HTTPException(status_code=400, detail="Проекты не выбраны")
    
    final = {}
    for project in request.projects:
        prompts_list = [
            prompts.GET_PROJECT_DETAILS_AIM.format(text=request.text, project=project),
            prompts.GET_PROJECT_DETAILS_VALUE.format(text=request.text, project=project),
            prompts.GET_PROJECT_DETAILS_BUDGET.format(text=request.text, project=project),
            prompts.GET_PROJECT_DETAILS_ECO_EFFECT.format(text=request.text, project=project),
            prompts.GET_PROJECT_DETAILS_DEADLINE.format(text=request.text, project=project),
            prompts.GET_PROJECT_DETAILS_NEW_PLAN.format(text=request.text, project=project),
            prompts.GET_PROJECT_DETAILS_CONCLUSION.format(text=request.text, project=project),
        ]
        final[project] = "\n\n".join([await generate_response(p) for p in prompts_list])
        final[project] = final[project].replace("Конец ответа", "").replace('</s>', '').strip()
    
    return {"result": final}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)