fix: Fix http empty headers error (#1710)

This commit is contained in:
Fangyin Cheng 2024-07-11 19:39:02 +08:00 committed by GitHub
parent 8c4c546cce
commit 125db534cb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 14 additions and 9 deletions

View File

@ -161,5 +161,6 @@ class RemoteModelWorker(ModelWorker):
def _get_trace_headers(self):
span_id = root_tracer.get_current_span_id()
headers = self.headers.copy()
headers.update({DBGPT_TRACER_SPAN_ID: span_id})
if span_id:
headers.update({DBGPT_TRACER_SPAN_ID: span_id})
return headers

View File

@ -693,9 +693,10 @@ class OpenAPIEmbeddings(BaseModel, Embeddings):
"""
# Call OpenAI Embedding API
headers = {}
if self.pass_trace_id:
current_span_id = root_tracer.get_current_span_id()
if self.pass_trace_id and current_span_id:
# Set the trace ID if available
headers[DBGPT_TRACER_SPAN_ID] = root_tracer.get_current_span_id()
headers[DBGPT_TRACER_SPAN_ID] = current_span_id
res = self.session.post( # type: ignore
self.api_url,
json={"input": texts, "model": self.model_name},
@ -726,9 +727,10 @@ class OpenAPIEmbeddings(BaseModel, Embeddings):
List[float] corresponds to a single input text.
"""
headers = {"Authorization": f"Bearer {self.api_key}"}
if self.pass_trace_id:
current_span_id = root_tracer.get_current_span_id()
if self.pass_trace_id and current_span_id:
# Set the trace ID if available
headers[DBGPT_TRACER_SPAN_ID] = root_tracer.get_current_span_id()
headers[DBGPT_TRACER_SPAN_ID] = current_span_id
async with aiohttp.ClientSession(
headers=headers, timeout=aiohttp.ClientTimeout(total=self.timeout)
) as session:

View File

@ -117,9 +117,10 @@ class OpenAPIRerankEmbeddings(BaseModel, RerankEmbeddings):
if not candidates:
return []
headers = {}
if self.pass_trace_id:
current_span_id = root_tracer.get_current_span_id()
if self.pass_trace_id and current_span_id:
# Set the trace ID if available
headers[DBGPT_TRACER_SPAN_ID] = root_tracer.get_current_span_id()
headers[DBGPT_TRACER_SPAN_ID] = current_span_id
data = {"model": self.model_name, "query": query, "documents": candidates}
response = self.session.post( # type: ignore
self.api_url, json=data, timeout=self.timeout, headers=headers
@ -130,9 +131,10 @@ class OpenAPIRerankEmbeddings(BaseModel, RerankEmbeddings):
async def apredict(self, query: str, candidates: List[str]) -> List[float]:
"""Predict the rank scores of the candidates asynchronously."""
headers = {"Authorization": f"Bearer {self.api_key}"}
if self.pass_trace_id:
current_span_id = root_tracer.get_current_span_id()
if self.pass_trace_id and current_span_id:
# Set the trace ID if available
headers[DBGPT_TRACER_SPAN_ID] = root_tracer.get_current_span_id()
headers[DBGPT_TRACER_SPAN_ID] = current_span_id
async with aiohttp.ClientSession(
headers=headers, timeout=aiohttp.ClientTimeout(total=self.timeout)
) as session: