Print logging

This commit is contained in:
Justin Hayes 2024-06-28 11:06:08 -04:00 committed by GitHub
parent c44217b5b8
commit b82d902185
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -2,15 +2,11 @@
from typing import List, Union, Iterator from typing import List, Union, Iterator
import os import os
import logging
from pydantic import BaseModel from pydantic import BaseModel
import google.generativeai as genai import google.generativeai as genai
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class Pipeline: class Pipeline:
"""Google GenAI pipeline""" """Google GenAI pipeline"""
@ -76,8 +72,8 @@ class Pipeline:
def pipe( def pipe(
self, user_message: str, model_id: str, messages: List[dict], body: dict self, user_message: str, model_id: str, messages: List[dict], body: dict
) -> Union[str, Iterator]: ) -> Union[str, Iterator]:
logger.info(f"Pipe function called for model: {model_id}") print(f"Pipe function called for model: {model_id}")
logger.info(f"Stream mode: {body['stream']}") print(f"Stream mode: {body['stream']}")
system_prompt = None system_prompt = None
google_messages = [] google_messages = []
@ -104,8 +100,8 @@ class Pipeline:
"parts": parts "parts": parts
}) })
except Exception as e: except Exception as e:
logger.error(f"Error processing message: {e}") print(f"Error processing message: {e}")
logger.error(f"Problematic message: {message}") print(f"Problematic message: {message}")
try: try:
model = genai.GenerativeModel( model = genai.GenerativeModel(
@ -124,16 +120,19 @@ class Pipeline:
) )
if body["stream"]: if body["stream"]:
logger.info("Streaming response") print("Streaming response")
for chunk in response: for chunk in response:
yield chunk.text yield chunk.text
return "" return ""
else: else:
logger.info("Non-streaming response") print("Non-streaming response")
result = response.text result = response.text
logger.info(f"Generated content: {result}") print(f"Generated content: {result}")
return result return result
except Exception as e: except Exception as e:
logger.error(f"Error generating content: {e}") print(f"Error generating content: {e}")
return f"An error occurred: {str(e)}" return f"An error occurred: {str(e)}"
finally:
print("Pipe function completed")