429 lines
17 KiB
Python
429 lines
17 KiB
Python
from config.settings import settings
|
|
from fastapi import APIRouter, Depends, HTTPException
|
|
from models.ai import (
|
|
AutoUpdateBibleRequest,
|
|
AutoUpdateBibleResponse,
|
|
CleanupDuplicatesResponse,
|
|
ExtractedEntities,
|
|
ExtractEntitiesRequest,
|
|
ExtractEntitiesResponse,
|
|
GenerateImageRequest,
|
|
GenerateImageResponse,
|
|
GenerateRequest,
|
|
GenerateResponse,
|
|
PreviewContextRequest,
|
|
PreviewContextResponse,
|
|
)
|
|
from models.bible import Character, Event, Location, StoryBible
|
|
from models.user import User
|
|
from services.bible_service import bible_service
|
|
from services.context_builder import context_builder
|
|
from services.deduplication import deduplication_service
|
|
from services.entity_extractor import entity_extractor
|
|
from services.image_generator import image_generator
|
|
from services.llm_router import llm_router
|
|
from services.model_config import model_config
|
|
|
|
from routers.auth import get_current_user
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def _get_user_project(project_id: str, current_user: User) -> None:
|
|
"""Проверить что проект принадлежит пользователю"""
|
|
from services.project_service import project_service
|
|
|
|
project = project_service.get_project(project_id)
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
if project.user_id != current_user.id:
|
|
raise HTTPException(status_code=403, detail="Access denied")
|
|
|
|
|
|
@router.get("/ai/providers/status")
|
|
async def get_providers_status():
|
|
"""Проверить, какие провайдеры настроены и готовы к использованию"""
|
|
return {
|
|
"providers": settings.validate(),
|
|
"openrouter_configured": bool(
|
|
settings.OPENROUTER_API_KEY
|
|
and settings.OPENROUTER_API_KEY != "your_api_key_here"
|
|
),
|
|
"openrouter_url": settings.OPENROUTER_BASE_URL,
|
|
"ollama_url": settings.OLLAMA_BASE_URL,
|
|
}
|
|
|
|
|
|
@router.get("/ai/models")
|
|
async def get_available_models(provider: str = "ollama"):
|
|
"""Получить список доступных моделей для указанного провайдера"""
|
|
try:
|
|
models = model_config.get_available_models(provider)
|
|
default_model = model_config.get_default_model(provider)
|
|
return {
|
|
"provider": provider,
|
|
"default_model": default_model,
|
|
"models": models,
|
|
}
|
|
except ValueError as e:
|
|
raise HTTPException(status_code=400, detail=str(e))
|
|
|
|
|
|
@router.post("/ai/generate", response_model=GenerateResponse)
|
|
async def generate_text(request: GenerateRequest):
|
|
try:
|
|
# Обогащаем промпт контекстом из Story Bible
|
|
enriched_prompt = context_builder.enrich_prompt(
|
|
original_prompt=request.prompt,
|
|
selected_text=request.selected_text or "",
|
|
bible=request.story_bible,
|
|
)
|
|
|
|
# Генерируем текст с обогащенным промптом
|
|
generated_text = await llm_router.generate(
|
|
prompt=enriched_prompt,
|
|
provider=request.provider,
|
|
model=request.model,
|
|
api_key=request.api_key,
|
|
preset=request.preset,
|
|
options=request.options,
|
|
)
|
|
return GenerateResponse(
|
|
generated_text=generated_text,
|
|
provider=request.provider,
|
|
model=request.model or "default",
|
|
preset=request.preset,
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.post("/ai/preview-context", response_model=PreviewContextResponse)
|
|
async def preview_context(request: PreviewContextRequest):
|
|
"""Предпросмотр контекста, который будет добавлен к промпту"""
|
|
mentioned = context_builder.extract_mentioned_entities(
|
|
text=request.text,
|
|
bible=request.story_bible,
|
|
)
|
|
|
|
context_str = context_builder.build_context_string(mentioned)
|
|
|
|
return PreviewContextResponse(
|
|
context_string=context_str,
|
|
mentioned_characters=[c.name for c in mentioned["characters"]],
|
|
mentioned_locations=[l.name for l in mentioned["locations"]],
|
|
mentioned_events=[e.title for e in mentioned["events"]],
|
|
)
|
|
|
|
|
|
@router.post("/ai/extract-entities", response_model=ExtractEntitiesResponse)
|
|
async def extract_entities(request: ExtractEntitiesRequest):
|
|
"""Извлечение сущностей из текста с помощью AI"""
|
|
try:
|
|
# Используем OpenRouter с моделью, оптимальной для извлечения
|
|
provider = "openrouter"
|
|
model = "qwen/qwen-2.5-7b-instruct"
|
|
|
|
if (
|
|
not settings.OPENROUTER_API_KEY
|
|
or settings.OPENROUTER_API_KEY == "your_api_key_here"
|
|
):
|
|
raise ValueError(
|
|
"OpenRouter API key not configured. "
|
|
"Entity extraction requires a cloud model. "
|
|
"Please set OPENROUTER_API_KEY in backend/.env"
|
|
)
|
|
|
|
entities = await entity_extractor.extract_entities(
|
|
text=request.text,
|
|
existing_bible=request.existing_bible,
|
|
provider=provider,
|
|
model=model,
|
|
)
|
|
|
|
return ExtractEntitiesResponse(
|
|
entities=ExtractedEntities(**entities),
|
|
model_used=model,
|
|
provider=provider,
|
|
)
|
|
except ValueError as e:
|
|
raise HTTPException(status_code=400, detail=str(e))
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.post("/ai/auto-update-bible", response_model=AutoUpdateBibleResponse)
|
|
async def auto_update_bible(
|
|
request: AutoUpdateBibleRequest, current_user: User = Depends(get_current_user)
|
|
):
|
|
"""Автоматическое обновление Story Bible при сохранении сцены"""
|
|
if not request.auto_update_enabled:
|
|
return AutoUpdateBibleResponse(updated=False, message="Auto-update disabled")
|
|
|
|
_get_user_project(request.project_id, current_user)
|
|
|
|
try:
|
|
# Получаем текущую Story Bible проекта
|
|
current_bible = bible_service.get_bible(request.project_id, current_user.id)
|
|
|
|
# Извлекаем сущности из текста сцены
|
|
entities = await entity_extractor.extract_entities(
|
|
text=request.scene_text,
|
|
existing_bible=current_bible,
|
|
provider="openrouter",
|
|
model="qwen/qwen-2.5-7b-instruct",
|
|
)
|
|
|
|
# Добавляем новые сущности в Bible
|
|
new_chars_count = 0
|
|
new_locs_count = 0
|
|
new_events_count = 0
|
|
|
|
# Добавляем новых персонажей С ДЕДУПЛИКАЦИЕЙ
|
|
for char_data in entities.get("new_characters", []):
|
|
try:
|
|
char_name = char_data.get("name", "")
|
|
if not char_name:
|
|
continue
|
|
|
|
# Ищем существующего персонажа с похожим именем
|
|
existing_char = deduplication_service.find_existing_character(
|
|
char_name, current_bible
|
|
)
|
|
|
|
if existing_char:
|
|
# Объединяем данные с существующим персонажем
|
|
deduplication_service.merge_character_data(existing_char, char_data)
|
|
print(f"Merged character: {char_name} -> {existing_char.name}")
|
|
else:
|
|
# Создаем нового персонажа
|
|
character = Character(**char_data)
|
|
current_bible.characters.append(character)
|
|
new_chars_count += 1
|
|
print(f"Added new character: {char_name}")
|
|
except Exception as e:
|
|
print(f"Failed to process character: {e}")
|
|
|
|
# Добавляем новые локации С ДЕДУПЛИКАЦИЕЙ
|
|
for loc_data in entities.get("new_locations", []):
|
|
try:
|
|
loc_name = loc_data.get("name", "")
|
|
if not loc_name:
|
|
continue
|
|
|
|
existing_loc = deduplication_service.find_existing_location(
|
|
loc_name, current_bible
|
|
)
|
|
|
|
if existing_loc:
|
|
# Обновляем описание, если новое длиннее
|
|
new_desc = loc_data.get("description", "")
|
|
if new_desc and len(new_desc) > len(existing_loc.description):
|
|
existing_loc.description = new_desc
|
|
|
|
new_atm = loc_data.get("atmosphere", "")
|
|
if new_atm and not existing_loc.atmosphere:
|
|
existing_loc.atmosphere = new_atm
|
|
|
|
print(f"Merged location: {loc_name} -> {existing_loc.name}")
|
|
else:
|
|
location = Location(**loc_data)
|
|
current_bible.locations.append(location)
|
|
new_locs_count += 1
|
|
print(f"Added new location: {loc_name}")
|
|
except Exception as e:
|
|
print(f"Failed to process location: {e}")
|
|
|
|
# Добавляем новые события С ДЕДУПЛИКАЦИЕЙ
|
|
for event_data in entities.get("new_events", []):
|
|
try:
|
|
event_title = event_data.get("title", "")
|
|
if not event_title:
|
|
continue
|
|
|
|
existing_event = deduplication_service.find_existing_event(
|
|
event_title, current_bible
|
|
)
|
|
|
|
if existing_event:
|
|
# Обновляем описание, если новое длиннее
|
|
new_desc = event_data.get("description", "")
|
|
if new_desc and len(new_desc) > len(existing_event.description):
|
|
existing_event.description = new_desc
|
|
print(f"Merged event: {event_title} -> {existing_event.title}")
|
|
else:
|
|
event = Event(**event_data)
|
|
current_bible.events.append(event)
|
|
new_events_count += 1
|
|
print(f"Added new event: {event_title}")
|
|
except Exception as e:
|
|
print(f"Failed to process event: {e}")
|
|
|
|
# Обновляем черты существующих персонажей
|
|
for trait_update in entities.get("updated_traits", []):
|
|
char_name = trait_update.get("character_name")
|
|
new_traits = trait_update.get("new_traits", [])
|
|
|
|
for char in current_bible.characters:
|
|
if char.name.lower() == (char_name or "").lower():
|
|
# Добавляем новые черты, которых еще нет
|
|
for trait in new_traits:
|
|
if trait not in char.traits:
|
|
char.traits.append(trait)
|
|
|
|
# Сохраняем обновленную Bible
|
|
if new_chars_count > 0 or new_locs_count > 0 or new_events_count > 0:
|
|
bible_service.update_bible(
|
|
request.project_id, current_bible, current_user.id
|
|
)
|
|
message = (
|
|
f"Story Bible обновлена: +{new_chars_count} персонажей, "
|
|
f"+{new_locs_count} локаций, +{new_events_count} событий"
|
|
)
|
|
else:
|
|
message = "Новых сущностей не обнаружено"
|
|
|
|
return AutoUpdateBibleResponse(
|
|
updated=True,
|
|
new_characters=new_chars_count,
|
|
new_locations=new_locs_count,
|
|
new_events=new_events_count,
|
|
message=message,
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Auto-update bible error: {e}")
|
|
return AutoUpdateBibleResponse(
|
|
updated=False,
|
|
message=f"Ошибка при обновлении: {str(e)}",
|
|
)
|
|
|
|
|
|
@router.post("/ai/extract-entities/test")
|
|
async def test_extraction():
|
|
"""Тестовый эндпоинт с примером текста"""
|
|
test_text = """
|
|
Алекс вошел в таверну "Серебряный дракон". Внутри пахло жареным мясом и элем.
|
|
За стойкой стоял хозяин — крепкий мужчина по имени Гаррет с рыжей бородой.
|
|
В углу сидела таинственная женщина в красном плаще — Алекс узнал Лиру.
|
|
"""
|
|
|
|
sample_bible = StoryBible(
|
|
characters=[
|
|
Character(name="Алекс", description="Главный герой", role="protagonist"),
|
|
Character(
|
|
name="Лира", description="Таинственная проводница", role="supporting"
|
|
),
|
|
]
|
|
)
|
|
|
|
try:
|
|
entities = await entity_extractor.extract_entities(
|
|
text=test_text,
|
|
existing_bible=sample_bible,
|
|
provider="openrouter",
|
|
model="qwen/qwen-2.5-7b-instruct",
|
|
)
|
|
|
|
return {
|
|
"test_text": test_text,
|
|
"extracted": entities,
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.post(
|
|
"/projects/{project_id}/bible/cleanup-duplicates",
|
|
response_model=CleanupDuplicatesResponse,
|
|
)
|
|
async def cleanup_duplicates(
|
|
project_id: str, current_user: User = Depends(get_current_user)
|
|
):
|
|
"""Очистка дубликатов в Story Bible"""
|
|
_get_user_project(project_id, current_user)
|
|
try:
|
|
bible = bible_service.get_bible(project_id, current_user.id)
|
|
|
|
# Удаляем дубликаты персонажей
|
|
unique_chars = []
|
|
removed_chars = 0
|
|
for char in bible.characters:
|
|
is_duplicate = False
|
|
for existing in unique_chars:
|
|
if deduplication_service.is_similar_name(char.name, existing.name):
|
|
# Объединяем данные
|
|
deduplication_service.merge_character_data(
|
|
existing, char.model_dump()
|
|
)
|
|
is_duplicate = True
|
|
removed_chars += 1
|
|
break
|
|
if not is_duplicate:
|
|
unique_chars.append(char)
|
|
bible.characters = unique_chars
|
|
|
|
# Удаляем дубликаты локаций
|
|
unique_locs = []
|
|
removed_locs = 0
|
|
for loc in bible.locations:
|
|
is_duplicate = False
|
|
for existing in unique_locs:
|
|
if deduplication_service.is_similar_name(loc.name, existing.name):
|
|
is_duplicate = True
|
|
removed_locs += 1
|
|
break
|
|
if not is_duplicate:
|
|
unique_locs.append(loc)
|
|
bible.locations = unique_locs
|
|
|
|
# Удаляем дубликаты событий
|
|
unique_events = []
|
|
removed_events = 0
|
|
for event in bible.events:
|
|
is_duplicate = False
|
|
for existing in unique_events:
|
|
if deduplication_service.is_similar_name(event.title, existing.title):
|
|
is_duplicate = True
|
|
removed_events += 1
|
|
break
|
|
if not is_duplicate:
|
|
unique_events.append(event)
|
|
bible.events = unique_events
|
|
|
|
# Сохраняем очищенную Bible
|
|
bible_service.update_bible(project_id, bible, current_user.id)
|
|
|
|
return CleanupDuplicatesResponse(
|
|
removed_characters=removed_chars,
|
|
removed_locations=removed_locs,
|
|
removed_events=removed_events,
|
|
message=f"Удалено дубликатов: {removed_chars} персонажей, {removed_locs} локаций, {removed_events} событий",
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.post("/ai/generate-image", response_model=GenerateImageResponse)
|
|
async def generate_image(
|
|
request: GenerateImageRequest, current_user: User = Depends(get_current_user)
|
|
):
|
|
"""Генерация изображения по промпту"""
|
|
try:
|
|
result = await image_generator.generate_image(
|
|
prompt=request.prompt,
|
|
project_id=request.project_id,
|
|
user_id=current_user.id,
|
|
document_id=request.document_id,
|
|
document_type=request.document_type,
|
|
model=request.model,
|
|
size=request.size,
|
|
)
|
|
return result
|
|
except ValueError as e:
|
|
raise HTTPException(status_code=400, detail=str(e))
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|