Merge pull request #92 from bukosabino/develop-add-langtrace

Adding langtrace to monitor the LLM and the embedding database traces

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

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
 

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