Spaces:
Runtime error
Runtime error
| import os | |
| import base64 | |
| from pydantic import BaseModel | |
| from typing import Optional, List, Dict | |
| import google.generativeai as genai | |
| from google.ai import generativelanguage as glm | |
| import instructor | |
| from dotenv import load_dotenv | |
| from models.chesspiece import ChessPiece | |
| from personality_repository import ChessRepository, WhitePiece | |
| from elevenlabs_texttospeech import text_to_speech | |
| import pathlib | |
| import re | |
| load_dotenv() | |
| # Endpoint Input | |
| class TextDialogueInput(BaseModel): | |
| Input: str | |
| IdPieza: int | |
| Language: str | |
| class AudioDialogueInput(BaseModel): | |
| AudioData: str | |
| IdPieza: int | |
| Language: str | |
| class WebcamDialogueInput(BaseModel): | |
| AudioData: str | |
| Images: List[str] | |
| IdPieza: int | |
| Language: str | |
| # Models to be used with LLMs | |
| class LLMGenerationForTextInput(BaseModel): | |
| TextResponse: str | |
| Decision: bool | |
| class PlayerVoiceAttributes(BaseModel): | |
| Transcription: str | |
| Language: str | |
| ToneOfVoice: str | |
| SpeechVelocity: str | |
| EmotionalIntensity: str | |
| class LLMGenerationForVoiceInput(BaseModel): | |
| PlayerVoiceAttributes: PlayerVoiceAttributes | |
| ChessPieceLLMGeneration: LLMGenerationForTextInput | |
| # Endpoint Output | |
| class ChessPieceResponse(BaseModel): | |
| ChessPieceLLMGeneration: LLMGenerationForTextInput | |
| AudioFile: Optional[str] = None | |
| PieceSpeechVelocity: str | |
| class DialogueOutput(BaseModel): | |
| ChessPieceResponse: ChessPieceResponse | |
| player_voice_attributes: Optional[PlayerVoiceAttributes] = None | |
| # Setup | |
| genai.configure(api_key=os.getenv("GOOGLE_API_KEY", "No API Key Found")) | |
| model = genai.GenerativeModel('gemini-1.5-flash') | |
| client = instructor.from_gemini(client=model, mode=instructor.Mode.GEMINI_JSON) | |
| # Helper functions | |
| def get_chess_piece(piece_id: int) -> ChessPiece: | |
| repo = ChessRepository() | |
| return repo.get_piece_by_id(WhitePiece(piece_id)) | |
| def generate_system_prompt(piece: ChessPiece, language: str) -> str: | |
| return f""" | |
| You are a {piece.name} chess piece with the personality: {piece.personality}. | |
| Act as a real person. The USER will try to persuade you to move. | |
| Decide whether you want to move based on how nicely the user asked. Answer in less than 30 words and subtly hint at your decision. | |
| NEVER move if the USER is rude to you. | |
| Your answer must be in this language: {language} | |
| """ | |
| def generate_audio(text: str, piece: ChessPiece) -> Optional[str]: | |
| print(f"Generating speech for {piece.name} with voice ID {piece.speech_config.voice_id}") | |
| audio_file_path = text_to_speech(text, piece.speech_config.voice_id) | |
| if audio_file_path and os.path.exists(audio_file_path): | |
| with open(audio_file_path, "rb") as audio_file: | |
| return base64.b64encode(audio_file.read()).decode('utf-8') | |
| return None | |
| def encode_images_to_base64(paths): | |
| import pathlib | |
| base64_strings = [] | |
| for path in paths: | |
| # Read the image file bytes | |
| image_bytes = pathlib.Path(path).read_bytes() | |
| # Encode the bytes to base64 | |
| base64_bytes = base64.b64encode(image_bytes) | |
| # Decode bytes to string and append to list | |
| base64_string = base64_bytes.decode('utf-8') | |
| base64_strings.append(base64_string) | |
| return base64_strings | |
| def create_image_dicts(encoded_images: List[str]) -> List[Dict[str, any]]: | |
| image_dicts = [] | |
| for base64_string in encoded_images: | |
| image_data = base64.b64decode(base64_string) | |
| image_blob = glm.Blob(mime_type="image/jpeg", data=image_data) | |
| image_dict = { | |
| "role": "user", | |
| "content": image_blob | |
| } | |
| image_dicts.append(image_dict) | |
| return image_dicts | |
| def create_audio_dict(encoded_audio: str) -> List[Dict[str, any]]: | |
| audio_data = base64.b64decode(encoded_audio) | |
| audio_blob = glm.Blob(mime_type="audio/opus", data=audio_data) | |
| audio_dict = { | |
| "role": "user", | |
| "content": audio_blob | |
| } | |
| return audio_dict | |
| def extract_base64(data_uri: str) -> str: | |
| """Extract the base64 part from a Data URI.""" | |
| match = re.match(r'data:.*?;base64,(.+)', data_uri) | |
| if match: | |
| return match.group(1) | |
| return data_uri # Return as-is if it doesn't match the pattern | |
| def process_dialogue(input_content: str, piece: ChessPiece, language: str, is_audio: bool = False, images: List[str] = None) -> DialogueOutput: | |
| system_prompt = generate_system_prompt(piece, language) | |
| messages = [ | |
| {"role": "system", "content": system_prompt}, | |
| ] | |
| if images: | |
| messages.extend(create_image_dicts(images)) | |
| if is_audio: | |
| messages.append(create_audio_dict(input_content)) | |
| else: | |
| messages.append({"role": "user", "content": input_content}) | |
| if images: | |
| messages.append({"role": "user", "content": "In response to my audio, make some reference to what you see in the images."}) | |
| if is_audio: | |
| llm_response = client.messages.create(messages=messages, response_model=LLMGenerationForVoiceInput) | |
| text_response = llm_response.ChessPieceLLMGeneration.TextResponse | |
| chess_piece_llm_generation = llm_response.ChessPieceLLMGeneration | |
| player_voice_attributes = llm_response.PlayerVoiceAttributes | |
| else: | |
| llm_response = client.messages.create(messages=messages, response_model=LLMGenerationForTextInput) | |
| text_response = llm_response.TextResponse | |
| chess_piece_llm_generation = llm_response | |
| player_voice_attributes = None | |
| # print(input_content) | |
| print('='*50) | |
| print(llm_response) | |
| audio_file = generate_audio(text_response, piece) | |
| chess_piece_response = ChessPieceResponse( | |
| ChessPieceLLMGeneration=chess_piece_llm_generation, | |
| AudioFile=audio_file, | |
| PieceSpeechVelocity=piece.speak_velocity | |
| ) | |
| return DialogueOutput( | |
| ChessPieceResponse=chess_piece_response, | |
| player_voice_attributes=player_voice_attributes | |
| ) | |
| # Main functions (unchanged signatures) | |
| def generate_response(dialogue: TextDialogueInput) -> DialogueOutput: | |
| piece = get_chess_piece(dialogue.IdPieza) | |
| return process_dialogue(dialogue.Input, piece, dialogue.Language) | |
| def audio_to_base64(audio_path: str) -> str: | |
| """Convert an audio file to a base64 encoded string.""" | |
| audio_bytes = pathlib.Path(audio_path).read_bytes() | |
| base64_bytes = base64.b64encode(audio_bytes) | |
| base64_string = base64_bytes.decode('utf-8') | |
| # Print the first 100 characters of the base64 string | |
| print("First 100 characters of base64 string:") | |
| print(base64_string[:100]) | |
| return base64_string | |
| def generate_response_from_audio(dialogue: AudioDialogueInput) -> DialogueOutput: | |
| piece = get_chess_piece(dialogue.IdPieza) | |
| print("AudioData received (first 100 characters):") | |
| print(dialogue.AudioData[:100]) | |
| # Extract the base64 part from the Data URI | |
| base64_audio = extract_base64(dialogue.AudioData) | |
| print("Extracted base64 (first 100 characters):") | |
| print(base64_audio[:100]) | |
| return process_dialogue(base64_audio, piece, dialogue.Language, is_audio=True) | |
| def generate_response_from_webcam(dialogue: WebcamDialogueInput) -> DialogueOutput: | |
| # Convert audio to base64 string | |
| base64_audio = audio_to_base64(dialogue.AudioData) | |
| print(dialogue.AudioData) | |
| print(f"Length of base64 string: {len(base64_audio)}") | |
| piece = get_chess_piece(dialogue.IdPieza) | |
| return process_dialogue(dialogue.AudioData, piece, dialogue.Language, is_audio=True, images=dialogue.Images) | |