modules/vertex_ai.py

import os
from typing import List, Optional
from google.cloud import aiplatform
from vertexai.language_models import TextEmbeddingModel
from vertexai.generative_models import GenerativeModel, Part

PROJECT_ID = os.getenv("GCP_PROJECT_ID")
LOCATION = os.getenv("GCP_LOCATION", "us-central1")

aiplatform.init(project=PROJECT_ID, location=LOCATION)

def generate_embedding(text: str) -> List[float]:
    """Generates a vector embedding."""
    model = TextEmbeddingModel.from_pretrained("text-embedding-004")
    embeddings = model.get_embeddings([text])
    return [float(v) for v in embeddings[0].values]

def synthesize_answer(query: str, context_chunks: List[str]) -> str:
    """Uses Gemini for grounded synthesis."""
    model = GenerativeModel("gemini-1.5-flash")
    context_text = "\n\n---\n\n".join(context_chunks)
    prompt = f"Answer the question using ONLY the context library.\nCONTEXT:\n{context_text}\nQUESTION:\n{query}"
    response = model.generate_content(prompt)
    return response.text

def generate_outline_from_context(topic: str, context_chunks: List[str]) -> str:
    """Uses context to draft a new book outline."""
    model = GenerativeModel("gemini-1.5-flash")
    context_text = "\n\n---\n\n".join(context_chunks)
    prompt = f"Based on the following knowledge library, draft a 10-chapter book outline for the topic: '{topic}'. Ensure it aligns with the existing knowledge style.\nCONTEXT:\n{context_text}"
    response = model.generate_content(prompt)
    return response.text