diff --git a/requirements.txt b/requirements.txt index 62515e8..94ddab4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,10 +19,9 @@ langsmith==0.1.65 langchain-openai==0.1.8 langchain-core==0.2.3 langchain-community==0.2.1 +langtrace_python_sdk==2.1.26 -qdrant-client==1.8.0 -supabase==1.0.2 -pinecone-client==2.2.2 +qdrant-client==1.9.2 sentence_transformers==2.2.2 openai==1.30.5 tavily-python==0.3.3 diff --git a/src/service/main.py b/src/service/main.py index 9996610..2e30733 100644 --- a/src/service/main.py +++ b/src/service/main.py @@ -2,13 +2,18 @@ import logging as lg import time import uuid +import os +import typing as tp import httpx -from fastapi import FastAPI +from fastapi import FastAPI, Request from src.initialize import initialize_app, initialize_logging -from src.utils import timeit +from src.utils import get_ip_client, inject_additional_attributes, timeit +from langtrace_python_sdk import SendUserFeedback, langtrace +from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span +langtrace.init(api_key=os.environ.get('LANGTRACE_API_KEY')) initialize_logging() APP = FastAPI() @@ -22,6 +27,18 @@ DEFAULT_COLLECTION_NAME = "justicio" +@with_langtrace_root_span() +async def call_llm_api(span_id, trace_id, model_name: str, messages: tp.List[tp.Dict[str, str]]): + response = await INIT_OBJECTS.openai_client.chat.completions.create( + model=model_name, + messages=messages, + temperature=INIT_OBJECTS.config_loader["temperature"], + seed=INIT_OBJECTS.config_loader["seed"], + max_tokens=INIT_OBJECTS.config_loader["max_tokens"], + ) + return response, span_id, trace_id + + @APP.get("/healthcheck") @timeit async def healthcheck(): @@ -67,9 +84,22 @@ async def a_request_get(url): return response.text +@APP.get("/qa_feedback") +@with_langtrace_root_span("Feedback") +@timeit +async def qa_feedback(span_id: str, trace_id: str, user_score: int): + data = { + "spanId": span_id, "traceId": trace_id, "userScore": user_score, "userId": None + } + SendUserFeedback().evaluate(data=data) + return {"feedback": "OK"} + + @APP.get("/qa") +@with_langtrace_root_span("RAG Justicio") @timeit async def qa( + request: Request, input_query: str = DEFAULT_INPUT_QUERY, collection_name: str = DEFAULT_COLLECTION_NAME, model_name: str = INIT_OBJECTS.config_loader["llm_model_name"], @@ -99,12 +129,9 @@ async def qa( {"role": "user", "content": input_query}, ] # logger.info(messages) - response = await INIT_OBJECTS.openai_client.chat.completions.create( - model=model_name, - messages=messages, - temperature=INIT_OBJECTS.config_loader["temperature"], - seed=INIT_OBJECTS.config_loader["seed"], - max_tokens=INIT_OBJECTS.config_loader["max_tokens"], + response, span_id, trace_id = await inject_additional_attributes( + lambda: call_llm_api(model_name=model_name, messages=messages), + {"db.collection.name": collection_name, "service.ip": get_ip_client(request)} ) answer = response.choices[0].message.content logger.info(answer) @@ -114,6 +141,8 @@ async def qa( scoring_id=str(uuid.uuid4()), context=docs, answer=answer, + span_id=str(span_id), + trace_id=str(trace_id), ) return response_payload diff --git a/src/utils.py b/src/utils.py index 199817c..69b2b8e 100644 --- a/src/utils.py +++ b/src/utils.py @@ -6,6 +6,11 @@ from langchain.schema import Document from langchain.vectorstores import SupabaseVectorStore from pydantic import BaseModel +from fastapi import Request +from opentelemetry import baggage, context +from langtrace_python_sdk.constants.instrumentation.common import ( + LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY, +) class StandardSupabaseVectorStore(SupabaseVectorStore): @@ -35,3 +40,22 @@ async def wrapper(*args, **kwargs): return result return wrapper + + +async def inject_additional_attributes(fn, attributes=None): + if attributes: + new_ctx = baggage.set_baggage( + LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY, attributes + ) + context.attach(new_ctx) + + return await fn() + + +def get_ip_client(request: Request): + x_forwarded_for = request.headers.get('x-forwarded-for') + if x_forwarded_for: + ip_client = x_forwarded_for.split(',')[0] + else: + ip_client = request.client.host + return ip_client