VietCat commited on
Commit
87b2780
·
1 Parent(s): 7595c71

refactor llm/embedding flow

Browse files
Files changed (2) hide show
  1. app/gemini_client.py +6 -2
  2. app/sheets.py +1 -1
app/gemini_client.py CHANGED
@@ -13,12 +13,16 @@ class GeminiClient:
13
  def generate_text(self, prompt: str, **kwargs) -> str:
14
  try:
15
  response = self._model.generate_content(prompt, **kwargs)
16
- logger.info(f"[GEMINI][RAW_RESPONSE] {response}")
 
17
  # response có thể là GenerativeModelResponse, lấy text hoặc trả về str
18
  if hasattr(response, 'text'):
 
19
  return response.text
20
  elif hasattr(response, 'candidates') and response.candidates:
 
21
  return response.candidates[0].content.parts[0].text
 
22
  return str(response)
23
  except Exception as e:
24
  logger.error(f"[GEMINI] Error: {e}")
@@ -38,7 +42,7 @@ class GeminiClient:
38
  content=text,
39
  task_type="retrieval_query"
40
  )
41
- logger.info(f"[GEMINI][EMBEDDING][RAW_RESPONSE] {response}")
42
  return response['embedding']
43
  except Exception as e:
44
  logger.error(f"[GEMINI][EMBEDDING] Error: {e}")
 
13
  def generate_text(self, prompt: str, **kwargs) -> str:
14
  try:
15
  response = self._model.generate_content(prompt, **kwargs)
16
+ if hasattr(response, 'usage_metadata'):
17
+ logger.info(f"[GEMINI][USAGE] Prompt Token Count: {response.usage_metadata.prompt_token_count} - Candidate Token Count: {response.usage_metadata.candidates_token_count} - Total Token Count: {response.usage_metadata.total_token_count}")
18
  # response có thể là GenerativeModelResponse, lấy text hoặc trả về str
19
  if hasattr(response, 'text'):
20
+ logger.info(f"[GEMINI][TEXT_RESPONSE] {response.text}")
21
  return response.text
22
  elif hasattr(response, 'candidates') and response.candidates:
23
+ logger.info(f"[GEMINI][CANDIDATES_RESPONSE] {response.candidates[0].content.parts[0].text}")
24
  return response.candidates[0].content.parts[0].text
25
+ logger.info(f"[GEMINI][RAW_RESPONSE] {response}")
26
  return str(response)
27
  except Exception as e:
28
  logger.error(f"[GEMINI] Error: {e}")
 
42
  content=text,
43
  task_type="retrieval_query"
44
  )
45
+ logger.info(f"[GEMINI][EMBEDDING][RAW_RESPONSE] {response['embedding'][:10]} ..... {response['embedding'][-10:]}")
46
  return response['embedding']
47
  except Exception as e:
48
  logger.error(f"[GEMINI][EMBEDDING] Error: {e}")
app/sheets.py CHANGED
@@ -146,7 +146,7 @@ class SheetsClient:
146
  range=SHEET_RANGE
147
  ).execute()
148
  values = result.get('values', [])
149
- logger.info(f"[DEBUG] Gsheet values {values}")
150
  ts = datetime.now().isoformat()
151
  # Đảm bảo timestamp luôn là list
152
  if timestamp is None:
 
146
  range=SHEET_RANGE
147
  ).execute()
148
  values = result.get('values', [])
149
+ # logger.info(f"[DEBUG] Gsheet values {values}")
150
  ts = datetime.now().isoformat()
151
  # Đảm bảo timestamp luôn là list
152
  if timestamp is None: