dlaima commited on
Commit
9edd6bb
·
verified ·
1 Parent(s): badada6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -19,8 +19,9 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
19
 
20
  # Wrapper class for model output
21
  class GenerationResult:
22
- def __init__(self, content):
23
  self.content = content
 
24
 
25
  # Gemini model wrapper
26
  class GeminiFlashModel:
@@ -45,9 +46,9 @@ class GeminiFlashModel:
45
 
46
  try:
47
  response = self.model.generate_content(prompt)
48
- return GenerationResult(response.text.strip())
49
  except Exception as e:
50
- return GenerationResult(f"GENERATION ERROR: {e}")
51
 
52
  # Agent using Gemini
53
  class MyAgent:
@@ -150,3 +151,4 @@ if __name__ == "__main__":
150
 
151
 
152
 
 
 
19
 
20
  # Wrapper class for model output
21
  class GenerationResult:
22
+ def __init__(self, content, token_usage=None):
23
  self.content = content
24
+ self.token_usage = token_usage or {}
25
 
26
  # Gemini model wrapper
27
  class GeminiFlashModel:
 
46
 
47
  try:
48
  response = self.model.generate_content(prompt)
49
+ return GenerationResult(content=response.text.strip(), token_usage={}) # Gemini doesn't expose usage
50
  except Exception as e:
51
+ return GenerationResult(content=f"GENERATION ERROR: {e}", token_usage={})
52
 
53
  # Agent using Gemini
54
  class MyAgent:
 
151
 
152
 
153
 
154
+