Skip to content

Commit 5285829

Browse files
CarltonXiangharvey_xiangfridayL
authored
Feat/add opentelmetry (#298)
* feat: custom_logger support info and error * feat: customLogger support info and error log * feat: replace all Thread to contextThread * feat: add debug logger info for nebular * feat: delete nebular * feat: debug logger info * revert: log.py console config * feat: remove useless log * fix: unit test error --------- Co-authored-by: harvey_xiang <[email protected]> Co-authored-by: chunyu li <[email protected]>
1 parent 709f87c commit 5285829

File tree

21 files changed

+301
-231
lines changed

21 files changed

+301
-231
lines changed

poetry.lock

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ dependencies = [
4242
"tenacity (>=9.1.2,<10.0.0)", # Error handling and retrying library
4343
"fastapi[all] (>=0.115.12,<0.116.0)", # Web framework for building APIs
4444
"sqlalchemy (>=2.0.41,<3.0.0)", # SQL toolkit
45+
"pymysql (>=1.1.0,<2.0.0)", # MySQL Python driver
4546
"scikit-learn (>=1.7.0,<2.0.0)", # Machine learning
4647
"fastmcp (>=2.10.5,<3.0.0)",
4748
"python-dateutil (>=2.9.0.post0,<3.0.0)",

src/memos/api/context/context_thread.py

Lines changed: 0 additions & 96 deletions
This file was deleted.
Lines changed: 15 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
import logging
22

3-
from fastapi import Depends, Header, Request
4-
5-
from memos.api.context.context import RequestContext, set_request_context
3+
from memos.context.context import RequestContext, get_current_context
64

75

86
logger = logging.getLogger(__name__)
@@ -11,56 +9,17 @@
119
G = RequestContext
1210

1311

14-
def get_trace_id_from_header(
15-
trace_id: str | None = Header(None, alias="trace-id"),
16-
x_trace_id: str | None = Header(None, alias="x-trace-id"),
17-
g_trace_id: str | None = Header(None, alias="g-trace-id"),
18-
) -> str | None:
19-
"""
20-
Extract trace_id from various possible headers.
21-
22-
Priority: g-trace-id > x-trace-id > trace-id
23-
"""
24-
return g_trace_id or x_trace_id or trace_id
25-
26-
27-
def get_request_context(
28-
request: Request, trace_id: str | None = Depends(get_trace_id_from_header)
29-
) -> RequestContext:
30-
"""
31-
Get request context object with trace_id and request metadata.
32-
33-
This function creates a RequestContext and automatically sets it
34-
in the global context for use throughout the request lifecycle.
35-
"""
36-
# Create context object
37-
ctx = RequestContext(trace_id=trace_id)
38-
39-
# Set the context globally for this request
40-
set_request_context(ctx)
41-
42-
# Log request start
43-
logger.info(f"Request started with trace_id: {ctx.trace_id}")
44-
45-
# Add request metadata to context
46-
ctx.set("method", request.method)
47-
ctx.set("path", request.url.path)
48-
ctx.set("client_ip", request.client.host if request.client else None)
49-
50-
return ctx
51-
52-
53-
def get_g_object(trace_id: str | None = Depends(get_trace_id_from_header)) -> G:
12+
def get_g_object() -> G:
5413
"""
5514
Get Flask g-like object for the current request.
56-
57-
This creates a RequestContext and sets it globally for access
58-
throughout the request lifecycle.
15+
Returns the context created by middleware.
5916
"""
60-
g = RequestContext(trace_id=trace_id)
61-
set_request_context(g)
62-
logger.info(f"Request g object created with trace_id: {g.trace_id}")
63-
return g
17+
ctx = get_current_context()
18+
if ctx is None:
19+
raise RuntimeError(
20+
"No request context available. Make sure RequestContextMiddleware is properly configured."
21+
)
22+
return ctx
6423

6524

6625
def get_current_g() -> G | None:
@@ -70,8 +29,6 @@ def get_current_g() -> G | None:
7029
Returns:
7130
The current request's g object if available, None otherwise.
7231
"""
73-
from memos.context import get_current_context
74-
7532
return get_current_context()
7633

7734

@@ -85,6 +42,9 @@ def require_g() -> G:
8542
Raises:
8643
RuntimeError: If called outside of a request context.
8744
"""
88-
from memos.context import require_context
89-
90-
return require_context()
45+
ctx = get_current_context()
46+
if ctx is None:
47+
raise RuntimeError(
48+
"No request context available. This function must be called within a request handler."
49+
)
50+
return ctx

src/memos/api/middleware/request_context.py

Lines changed: 6 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,26 +10,17 @@
1010
from starlette.requests import Request
1111
from starlette.responses import Response
1212

13-
from memos.api.context.context import RequestContext, generate_trace_id, set_request_context
13+
from memos.context.context import RequestContext, generate_trace_id, set_request_context
1414

1515

1616
logger = logging.getLogger(__name__)
1717

1818

1919
def extract_trace_id_from_headers(request: Request) -> str | None:
2020
"""Extract trace_id from various possible headers with priority: g-trace-id > x-trace-id > trace-id."""
21-
trace_id = request.headers.get("g-trace-id")
22-
if trace_id:
23-
return trace_id
24-
25-
trace_id = request.headers.get("x-trace-id")
26-
if trace_id:
27-
return trace_id
28-
29-
trace_id = request.headers.get("trace-id")
30-
if trace_id:
31-
return trace_id
32-
21+
for header in ["g-trace-id", "x-trace-id", "trace-id"]:
22+
if trace_id := request.headers.get(header):
23+
return trace_id
3324
return None
3425

3526

@@ -45,19 +36,12 @@ class RequestContextMiddleware(BaseHTTPMiddleware):
4536

4637
async def dispatch(self, request: Request, call_next: Callable) -> Response:
4738
# Extract or generate trace_id
48-
trace_id = extract_trace_id_from_headers(request)
49-
if not trace_id:
50-
trace_id = generate_trace_id()
39+
trace_id = extract_trace_id_from_headers(request) or generate_trace_id()
5140

5241
# Create and set request context
53-
context = RequestContext(trace_id=trace_id)
42+
context = RequestContext(trace_id=trace_id, api_path=request.url.path)
5443
set_request_context(context)
5544

56-
# Add request metadata to context
57-
context.set("method", request.method)
58-
context.set("path", request.url.path)
59-
context.set("client_ip", request.client.host if request.client else None)
60-
6145
# Log request start with parameters
6246
params_log = {}
6347

0 commit comments

Comments
 (0)