Spaces:
Sleeping
Sleeping
File size: 10,906 Bytes
a8c7066 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | import time
import os
import uuid
from src.application.dto.ask_question_dto import AskQuestionCommand, AskQuestionResult, ChunkView
from src.application.ports.outbound.evaluation_port import EvaluationPort
from src.application.ports.inbound.ask_question_use_case import AskQuestionUseCase
from src.application.ports.outbound.llm_port import LlmPort
from src.application.ports.outbound.observability_port import ObservabilityPort, ScoreInput
from src.application.ports.outbound.retrieval_port import RetrievalPort
from src.domain.errors.exceptions import EmptyQuestionError
from src.domain.model.evaluation_models import AnswerEvaluation
from src.domain.model.rag_models import GenerationResult
MODEL_PRICING_PER_1K = {
"gemini-2.5-flash": {"input": 0.0003, "output": 0.0025},
"Qwen/Qwen2.5-3B-Instruct": {"input": 0.0, "output": 0.0},
}
class AskQuestionService(AskQuestionUseCase):
def __init__(
self,
retrieval_port: RetrievalPort,
llm_port: LlmPort,
observability_port: ObservabilityPort,
default_answer: str,
evaluation_port: EvaluationPort | None = None,
) -> None:
self._retrieval = retrieval_port
self._llm = llm_port
self._obs = observability_port
self._default_answer = default_answer
self._evaluation = evaluation_port
def execute(self, command: AskQuestionCommand) -> AskQuestionResult:
if not command.question or not command.question.strip():
raise EmptyQuestionError("Question is empty")
lambda_mult = command.lambda_mult if command.use_mmr else None
start_total = time.perf_counter()
model_name = self._llm.get_model_name()
pricing = MODEL_PRICING_PER_1K.get(model_name or "", {"input": 0.0, "output": 0.0})
# A new session_id is generated for every API call: the backend is stateless,
# unlike app.py which keeps one session_id per Gradio conversation.
session_id = str(uuid.uuid4())
trace_handle = self._obs.start_trace(
name="rag-pipeline",
payload={
"query": command.question,
"top_k": command.top_k,
"use_mmr": command.use_mmr,
"lambda_mult": lambda_mult,
"similarity_threshold": command.similarity_threshold,
},
trace_name=command.trace_name,
user_id=command.user_id,
session_id=session_id,
)
retrieval_handle = self._obs.start_observation(name="retrieval", as_type="retriever")
generation_handle = None
try:
retrieval_start = time.perf_counter()
list_chunks = self._retrieval.retrieve(
question=command.question,
top_k=command.top_k,
lambda_mult=lambda_mult,
similarity_threshold=command.similarity_threshold,
)
retrieval_latency = time.perf_counter() - retrieval_start
max_len_chunk = max((len(chunk.content) for chunk in list_chunks), default=0)
total_doc_length = sum(len(chunk.content) for chunk in list_chunks)
if retrieval_handle is not None:
retrieval_handle.update(
input={"query": command.question},
output={"retrieved_chunks": [chunk.content for chunk in list_chunks]},
status_message="SUCCESS",
metadata={
"reranker_used": False,
"mmr_used": command.use_mmr,
"top_k": str(command.top_k),
"mmr_lambda": str(lambda_mult) if command.use_mmr else None,
"reranker_name": None,
"similarity_threshold": str(command.similarity_threshold) if command.similarity_threshold is not None else None,
"retrieval_api_url": os.getenv("API_URL"),
},
)
retrieval_scores = {
"retrieve_latency": ScoreInput(name="retrieve_latency", value=round(retrieval_latency, 2)),
"max_len_chunk": ScoreInput(name="max_len_chunk", value=int(max_len_chunk)),
"total_doc_length": ScoreInput(name="total_doc_length", value=int(total_doc_length)),
}
retrieval_handle.score(retrieval_scores)
retrieval_handle.close()
retrieval_handle = None # For the finally block
except Exception as exc:
if retrieval_handle is not None:
retrieval_handle.update(input={"query": command.question}, output=str(exc), status_message="ERROR")
if trace_handle is not None:
trace_handle.update(input={"query": command.question}, output=str(exc), status_message="ERROR")
raise
try:
generation_handle = self._obs.start_observation(name="answer-generation", as_type="generation", model=model_name)
if not list_chunks:
answer = self._default_answer
generation_result = GenerationResult(
text=answer,
prompt="",
input_tokens=0,
output_tokens=0,
thoughts_tokens=0,
finish_reason=None,
)
if generation_handle is not None:
generation_handle.update(
input={"query": command.question},
output={"answer": answer},
status_message="DEFAULT_ANSWER",
)
if trace_handle is not None:
trace_handle.update(
input={"query": command.question},
output={"answer": answer},
status_message="DEFAULT_ANSWER")
else:
generation_start = time.perf_counter()
generation_result = self._llm.generate_answer(question=command.question, chunks=list_chunks)
generation_latency = time.perf_counter() - generation_start
answer = generation_result.text
if generation_handle is not None:
generation_config = self._llm.get_generation_config()
generation_handle.update(
input={"prompt": generation_result.prompt},
output={"answer": answer},
status_message="SUCCESS",
metadata={
"finish_reason": generation_result.finish_reason,
"prompt_version_file": "prompt_rag",
"thinking_budget": generation_config.get("thinking_budget"),
"max_output_tokens": generation_config.get("max_output_tokens"),
},
usage_details={
"input_tokens": int(generation_result.input_tokens),
"output_tokens": int(generation_result.output_tokens),
"thoughts_tokens": int(generation_result.thoughts_tokens),
},
)
if trace_handle is not None:
trace_handle.update(
input={"query": command.question},
output={"answer": answer},
status_message="SUCCESS",
)
output_tokens_total = generation_result.output_tokens + generation_result.thoughts_tokens
tokens_per_sec = int(output_tokens_total / generation_latency) if generation_latency > 0 else 0
input_cost = (generation_result.input_tokens / 1000.0) * pricing["input"]
output_cost = (output_tokens_total / 1000.0) * pricing["output"]
if generation_handle is not None:
generation_scores = {
"model_latency": ScoreInput(name="model_latency", value=round(generation_latency, 2), data_type="NUMERIC"),
"tokens_per_sec": ScoreInput(name="tokens_per_sec", value=int(tokens_per_sec), data_type="NUMERIC"),
"input_tokens": ScoreInput(name="input_tokens", value=int(generation_result.input_tokens), data_type="NUMERIC"),
"output_tokens": ScoreInput(name="output_tokens", value=int(output_tokens_total), data_type="NUMERIC"),
"total_cost": ScoreInput(name="total_cost", value=float(input_cost + output_cost), data_type="NUMERIC"),
}
generation_handle.score(generation_scores)
except Exception as exc:
if generation_handle is not None:
generation_handle.update(input={"query": command.question}, output=str(exc), status_message="ERROR")
if trace_handle is not None:
trace_handle.update(input={"query": command.question}, output=str(exc), status_message="ERROR")
raise
try:
total_latency = time.perf_counter() - start_total
if trace_handle is not None:
trace_handle.score({"total_latency": ScoreInput(name="total_latency", value=float(round(total_latency, 2)))})
if self._evaluation is not None:
self._evaluation.enqueue(
AnswerEvaluation(
question=command.question,
answer=answer,
contexts=[chunk.content for chunk in list_chunks],
),
trace_id=trace_handle.trace_id if trace_handle else None,
observation_id=getattr(generation_handle, "observation_id", None),
)
return AskQuestionResult(
answer=answer,
chunks=[
ChunkView(
id=chunk.id,
content=chunk.content,
score=chunk.score,
title=chunk.title,
author=chunk.author,
)
for chunk in list_chunks
],
trace_id=trace_handle.trace_id if trace_handle else None,
)
except Exception as exc:
if trace_handle is not None:
trace_handle.update(input={"query": command.question}, output=str(exc), status_message="ERROR")
raise
finally:
if generation_handle is not None:
generation_handle.close()
if retrieval_handle is not None:
retrieval_handle.close()
if trace_handle is not None:
trace_handle.close()
self._obs.finish_trace()
|