Spaces:

IOE-IFAD
/

test

Paused

Automated build & deploy

e454383 verified about 1 year ago

4.83 kB

import os

from dotenv import load_dotenv
from typing import List, Optional, Dict
from pydantic import BaseModel
from fastapi import Body
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from utils.ai import AiService
from utils.vector_db import VectorDB
from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import gpt_4_1_mini_complete, openai_embed
from utils.time_funcs import extract_timestamps
from utils.document_management import load_document, transcript_chunker, merge_chunks

from prompts import name_generation_prompt

# Load environment variables
load_dotenv()

# Initialize FastAPI app
app = FastAPI(title="Transcription API")

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Initialize services
working_dir = os.path.join(os.getcwd(), "_backend", "data","honduras_data")
if os.getenv("ENVIRONMENT") == "production":
    working_dir = os.path.join(os.getcwd(), "data","honduras_data")

transcription_data_dir = os.path.join(working_dir)
vector_db_name = "transcription"
vector_db_embedding_length = 1536

vector_db = VectorDB(
    vector_db_name,
    vector_db_embedding_length,
    os.getenv("OPENAI_API_KEY"),
    working_dir
)

ai_service = AiService(model="gpt-4.1")

lightrag_dir = os.path.join(working_dir, "lightrag")
rag = LightRAG(
    working_dir=lightrag_dir,
    llm_model_func=gpt_4_1_mini_complete,
    embedding_func=openai_embed
)

class ChatRequest(BaseModel):
    message: str
    chat_history: List[Dict[str, str]]


class ChatResponse(BaseModel):
    response: str
    timestamps: List[Dict[str, str | int]]  # Updated type annotation


class NameRequest(BaseModel):
    user_query: str
    assistant_response: str


class SearchRequest(BaseModel):
    query: str
    filters: Optional[Dict] = {}


class SessionInfoResponse(BaseModel):
    session_info: Dict
    video_url: str


# API Endpoints
@app.post("/api/v1/chat", response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest):
    try:
        
        translated_query = ai_service.generate_single_response(
            request.message,
            "Translate the following text to Spanish",
            json_output=False
        )
        # Search documents
        search_docs = vector_db.search(translated_query, limit=100,
                                        # filter=filters,
                                        score_threshold=0
                                        )

        # Get GraphML response using asyncio.create_task to avoid event loop issues
        import asyncio
        loop = asyncio.get_event_loop()
        graphML_resp = await loop.run_in_executor(
            None,
            lambda: rag.query(
                request.message,
                param=QueryParam(
                    mode="hybrid",
                )
            )
        )

        # Generate response
        assistant_response = ai_service.generate_response(
            request.chat_history,
            docs=search_docs,
            extra_context=graphML_resp,
        )

        # Extract timestamps
        timestamps = extract_timestamps(assistant_response, search_docs)

        return ChatResponse(
            response=assistant_response,
            timestamps=timestamps  # Remove json.dumps(), return list directly
        )
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/v1/generate_name", response_model=str)
async def generate_name_endpoint(request: NameRequest):
    try:
        messages = ai_service.init_messages(name_generation_prompt)
        ai_service.add_user_message(messages, request.user_query)

        # Generate response
        response = ai_service.generate_response(messages)

        return response
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/v1/transcript", response_model=List[Dict])
async def get_transcript_file(body: Dict = Body(...)):
    try:
        file_name = body.get("file_name")
        if not file_name:
            raise HTTPException(status_code=400, detail="file_name is required")
        file_path = os.path.join(transcription_data_dir, file_name)
        doc = load_document(file_path)
        transcript = transcript_chunker(doc, filename=file_name)
        transcript = merge_chunks(transcript)
        return transcript

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="File not found")


@app.get("/api/v1/health")
async def health_check():
    return {"status": "healthy"}

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

	import os

	from dotenv import load_dotenv
	from typing import List, Optional, Dict
	from pydantic import BaseModel
	from fastapi import Body
	from fastapi import FastAPI, HTTPException
	from fastapi.middleware.cors import CORSMiddleware
	from utils.ai import AiService
	from utils.vector_db import VectorDB
	from lightrag import LightRAG, QueryParam
	from lightrag.llm.openai import gpt_4_1_mini_complete, openai_embed
	from utils.time_funcs import extract_timestamps
	from utils.document_management import load_document, transcript_chunker, merge_chunks

	from prompts import name_generation_prompt

	# Load environment variables
	load_dotenv()

	# Initialize FastAPI app
	app = FastAPI(title="Transcription API")

	# Add CORS middleware
	app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"],
	allow_credentials=True,
	allow_methods=["*"], # Allows all methods
	allow_headers=["*"], # Allows all headers
	)

	# Initialize services
	working_dir = os.path.join(os.getcwd(), "_backend", "data","honduras_data")
	if os.getenv("ENVIRONMENT") == "production":
	working_dir = os.path.join(os.getcwd(), "data","honduras_data")

	transcription_data_dir = os.path.join(working_dir)
	vector_db_name = "transcription"
	vector_db_embedding_length = 1536

	vector_db = VectorDB(
	vector_db_name,
	vector_db_embedding_length,
	os.getenv("OPENAI_API_KEY"),
	working_dir
	)

	ai_service = AiService(model="gpt-4.1")

	lightrag_dir = os.path.join(working_dir, "lightrag")
	rag = LightRAG(
	working_dir=lightrag_dir,
	llm_model_func=gpt_4_1_mini_complete,
	embedding_func=openai_embed
	)

	class ChatRequest(BaseModel):
	message: str
	chat_history: List[Dict[str, str]]


	class ChatResponse(BaseModel):
	response: str
	timestamps: List[Dict[str, str \| int]] # Updated type annotation


	class NameRequest(BaseModel):
	user_query: str
	assistant_response: str


	class SearchRequest(BaseModel):
	query: str
	filters: Optional[Dict] = {}


	class SessionInfoResponse(BaseModel):
	session_info: Dict
	video_url: str


	# API Endpoints
	@app.post("/api/v1/chat", response_model=ChatResponse)
	async def chat_endpoint(request: ChatRequest):
	try:

	translated_query = ai_service.generate_single_response(
	request.message,
	"Translate the following text to Spanish",
	json_output=False
	)
	# Search documents
	search_docs = vector_db.search(translated_query, limit=100,
	# filter=filters,
	score_threshold=0
	)

	# Get GraphML response using asyncio.create_task to avoid event loop issues
	import asyncio
	loop = asyncio.get_event_loop()
	graphML_resp = await loop.run_in_executor(
	None,
	lambda: rag.query(
	request.message,
	param=QueryParam(
	mode="hybrid",
	)
	)
	)

	# Generate response
	assistant_response = ai_service.generate_response(
	request.chat_history,
	docs=search_docs,
	extra_context=graphML_resp,
	)

	# Extract timestamps
	timestamps = extract_timestamps(assistant_response, search_docs)

	return ChatResponse(
	response=assistant_response,
	timestamps=timestamps # Remove json.dumps(), return list directly
	)
	except Exception as e:
	print(e)
	raise HTTPException(status_code=500, detail=str(e))


	@app.post("/api/v1/generate_name", response_model=str)
	async def generate_name_endpoint(request: NameRequest):
	try:
	messages = ai_service.init_messages(name_generation_prompt)
	ai_service.add_user_message(messages, request.user_query)

	# Generate response
	response = ai_service.generate_response(messages)

	return response
	except Exception as e:
	raise HTTPException(status_code=500, detail=str(e))


	@app.post("/api/v1/transcript", response_model=List[Dict])
	async def get_transcript_file(body: Dict = Body(...)):
	try:
	file_name = body.get("file_name")
	if not file_name:
	raise HTTPException(status_code=400, detail="file_name is required")
	file_path = os.path.join(transcription_data_dir, file_name)
	doc = load_document(file_path)
	transcript = transcript_chunker(doc, filename=file_name)
	transcript = merge_chunks(transcript)
	return transcript

	except FileNotFoundError:
	raise HTTPException(status_code=404, detail="File not found")


	@app.get("/api/v1/health")
	async def health_check():
	return {"status": "healthy"}

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