154 lines
6.4 KiB
Python
154 lines
6.4 KiB
Python
# D:\GNPL-Private\TimeCraft\backend\app\services\llm_service.py
|
|
|
|
import os
|
|
import json
|
|
import httpx
|
|
from typing import List
|
|
from datetime import date
|
|
from dotenv import load_dotenv
|
|
from app.models.task import TaskCreateRequest, ScheduleResponse
|
|
|
|
load_dotenv()
|
|
|
|
LLM_PROVIDER = os.getenv("LLM_PROVIDER", "OLLAMA")
|
|
LLM_API_KEY = os.getenv("LLM_API_KEY")
|
|
LLM_API_URL = os.getenv("LLM_API_URL")
|
|
HF_MODEL_NAME = os.getenv("HF_MODEL_NAME", "google/gemma-2b-it")
|
|
OLLAMA_MODEL_NAME = os.getenv("OLLAMA_MODEL_NAME", "deepseek-r1:1.5b")
|
|
OLLAMA_API_URL = os.getenv("OLLAMA_API_URL", "http://localhost:11434/api/generate")
|
|
|
|
tokenizer = None
|
|
model = None
|
|
|
|
if LLM_PROVIDER == "HUGGINGFACE":
|
|
try:
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
import torch
|
|
tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)
|
|
model = AutoModelForCausalLM.from_pretrained(HF_MODEL_NAME, torch_dtype=torch.float16)
|
|
if torch.cuda.is_available():
|
|
model.to('cuda')
|
|
print(f"[LLM] Hugging Face model '{HF_MODEL_NAME}' loaded.")
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to load Hugging Face model: {e}")
|
|
raise RuntimeError("Hugging Face model loading failed.")
|
|
|
|
def infer_priority(task: TaskCreateRequest) -> str:
|
|
desc = task.task_description.lower()
|
|
if any(k in desc for k in ["non-negotiable", "critical", "mandatory", "must", "cannot be moved", "very high priority"]):
|
|
return "high"
|
|
if any(k in desc for k in ["optional", "nice to have", "if possible", "fun"]):
|
|
return "low"
|
|
return "medium"
|
|
|
|
async def get_scheduled_tasks(tasks: List[TaskCreateRequest]) -> ScheduleResponse:
|
|
today_str = date.today().isoformat()
|
|
prompt_tasks_data = [
|
|
{
|
|
**json.loads(task.model_dump_json()),
|
|
"original_task_id": f"task{index + 1}",
|
|
"priority": infer_priority(task)
|
|
}
|
|
for index, task in enumerate(tasks)
|
|
]
|
|
|
|
system_message = (
|
|
f"You are an AI-powered schedule and task planning assistant. Your job is to create a conflict-free, "
|
|
f"intelligent 30-day schedule from today ({today_str}).\n\n"
|
|
"Parse task details precisely: type, time_settings (start_time, end_time, duration), category, and priority.\n"
|
|
"Avoid time conflicts. If two fixed-time tasks clash, prioritize the 'high' priority one. Try to shift "
|
|
"'medium' priority tasks to nearby times. Skip low priority if no space.\n\n"
|
|
"Respect fixed times first. For flexible tasks, use hints.\n"
|
|
"Return strictly valid JSON:\n"
|
|
"{\n"
|
|
" \"scheduled_tasks\": [\n"
|
|
" {\n"
|
|
" \"task_name\": \"string\",\n"
|
|
" \"date\": \"YYYY-MM-DD\",\n"
|
|
" \"start_time\": \"HH:MM\",\n"
|
|
" \"end_time\": \"HH:MM\",\n"
|
|
" \"category\": \"personal | work | errands | extra\",\n"
|
|
" \"original_task_id\": \"string\"\n"
|
|
" }\n"
|
|
" ]\n"
|
|
"}"
|
|
)
|
|
|
|
user_message = f"Here are the tasks to schedule:\n{json.dumps(prompt_tasks_data, indent=2)}"
|
|
full_prompt = f"{system_message}\n\n{user_message}"
|
|
llm_output_content = ""
|
|
|
|
try:
|
|
if LLM_PROVIDER == "GENERIC_API":
|
|
if not LLM_API_URL or not LLM_API_KEY:
|
|
raise ValueError("Missing LLM_API_URL or LLM_API_KEY in .env")
|
|
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {LLM_API_KEY}"
|
|
}
|
|
|
|
payload = {
|
|
"model": "your-preferred-llm-model-name",
|
|
"messages": [
|
|
{"role": "system", "content": system_message},
|
|
{"role": "user", "content": user_message}
|
|
],
|
|
"response_format": {"type": "json_object"}
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=60.0)) as client:
|
|
response = await client.post(LLM_API_URL, headers=headers, json=payload)
|
|
response.raise_for_status()
|
|
llm_output_content = response.json().get("choices", [{}])[0].get("message", {}).get("content", "")
|
|
|
|
elif LLM_PROVIDER == "HUGGINGFACE":
|
|
import torch
|
|
input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids
|
|
if torch.cuda.is_available():
|
|
input_ids = input_ids.to("cuda")
|
|
outputs = model.generate(input_ids, max_new_tokens=1024, pad_token_id=tokenizer.eos_token_id)
|
|
llm_output_content = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
|
elif LLM_PROVIDER == "OLLAMA":
|
|
payload = {
|
|
"model": OLLAMA_MODEL_NAME,
|
|
"prompt": full_prompt,
|
|
"stream": False,
|
|
"format": "json"
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=60.0)) as client:
|
|
print("[LLM] Sending prompt to Ollama...")
|
|
response = await client.post(OLLAMA_API_URL, json=payload)
|
|
response.raise_for_status()
|
|
llm_output_content = response.json().get("response", "")
|
|
|
|
print("[LLM] Received response (first 500 chars):")
|
|
print(llm_output_content[:500])
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported LLM_PROVIDER: {LLM_PROVIDER}")
|
|
|
|
json_start = llm_output_content.find('{')
|
|
json_end = llm_output_content.rfind('}')
|
|
if json_start == -1 or json_end == -1:
|
|
raise ValueError("No valid JSON object found in LLM output.")
|
|
json_string = llm_output_content[json_start:json_end + 1]
|
|
parsed_data = json.loads(json_string)
|
|
return ScheduleResponse(**parsed_data)
|
|
|
|
except httpx.HTTPStatusError as e:
|
|
print(f"[ERROR] Ollama returned HTTP {e.response.status_code}: {e.response.text}")
|
|
raise RuntimeError(f"LLM API returned error {e.response.status_code}: {e.response.text}")
|
|
except httpx.RequestError as e:
|
|
print(f"[ERROR] Request to LLM service failed: {e}")
|
|
raise RuntimeError(f"Failed to connect to LLM service. Details: {e}")
|
|
except json.JSONDecodeError as e:
|
|
print(f"[ERROR] JSON decoding failed: {e}")
|
|
print(f"[DEBUG] Raw output: {llm_output_content}")
|
|
raise ValueError("LLM did not return valid JSON.")
|
|
except Exception as e:
|
|
print(f"[ERROR] Unexpected exception: {e}")
|
|
raise RuntimeError(f"Failed to schedule tasks due to an issue with the LLM service. Details: {e}")
|