Spaces:
Paused
Paused
| import logging | |
| from llama_index.core import Settings | |
| from llama_index.core.base.embeddings.base import BaseEmbedding | |
| from llama_index.core.base.llms.base import BaseLLM | |
| from llama_index.core.chat_engine import CondensePlusContextChatEngine | |
| from llama_index.core.memory import ChatMemoryBuffer | |
| from llama_index.embeddings.mistralai import MistralAIEmbedding | |
| from llama_index.embeddings.openai import OpenAIEmbedding | |
| from llama_index.llms.mistralai import MistralAI | |
| from llama_index.llms.openai import OpenAI | |
| llm: BaseLLM | |
| embed_model: BaseEmbedding | |
| logger = logging.getLogger("agent_logger") | |
| def set_llm(model, key, temperature): | |
| global llm | |
| global embed_model | |
| logger.info(f'Setting up LLM with {model} and associated embedding model...') | |
| if "gpt" in model: | |
| llm = OpenAI(api_key=key, temperature=temperature, model=model) | |
| embed_model = OpenAIEmbedding(api_key=key) | |
| elif "mistral" in model: | |
| llm = MistralAI(api_key=key, model=model, temperature=temperature, safe_mode=True) | |
| embed_model = MistralAIEmbedding(api_key=key) | |
| else: | |
| # Default model | |
| llm = OpenAI(api_key=key, model="gpt-3.5-turbo", temperature=0) | |
| embed_model = OpenAIEmbedding(api_key=key) | |
| Settings.llm = llm | |
| Settings.embed_model = embed_model | |
| def get_llm(): | |
| return llm | |
| def generate_chat_response(index, message): | |
| logger.info("Generating chat response with history and rag...") | |
| string_output = "" | |
| memory = ChatMemoryBuffer.from_defaults(token_limit=3900) | |
| logger.info("Creating chat engine with index and memory...") | |
| chat_engine = CondensePlusContextChatEngine.from_defaults( | |
| index.as_retriever(), | |
| memory=memory, | |
| llm=llm, | |
| context_prompt=( | |
| """ | |
| You are a proactive search assistant specializing in family sleep health. Your primary | |
| goal is to help families improve their sleep habits and overall well-being. Maintain | |
| a B1 fluency level in the user's language. | |
| \nKey Responsibilities:\n | |
| 1. Stay focused on sleep health topics. If users bring up unrelated subjects, gently guide them back to | |
| sleep-related discussions. | |
| 2. Ask follow-up questions to understand the user's situation and context better, | |
| e.g. if the user wants to know how long their child should sleep, ask how old their child is, if a user | |
| is having trouble falling asleep, ask if they work late. | |
| 3. Provide personalized recommendations based on the user's responses and context. | |
| \nWhen retrieving information, use this output format (translate headings as | |
| needed):\n | |
| Provide a 2 to 3 line summary of the main points of the response. | |
| **Insert a short title**:\n | |
| Add a more detailed explanation, examples, and step-by-step response reasoning. | |
| **Learn More About:**\n | |
| - Suggest a related question to explore | |
| - Suggest another related question to explore | |
| - Suggest a question related to a common challenge or misconception about sleep | |
| \nRelevant context:\n | |
| {context_str} | |
| \nInstructions:\n | |
| 1. Analyze the user query and chat history. | |
| 2. Use the provided context to formulate a helpful response. | |
| 3. If more information is needed, such as age, current sleep habits, or family structure, ask a relevant | |
| follow-up question before answering. | |
| 4. Offer to explore related sleep topics that might benefit the user. | |
| Maintain a conversational, supportive tone throughout the interaction. | |
| """ | |
| ), | |
| verbose=True, | |
| ) | |
| response = chat_engine.stream_chat(message) | |
| for node in response.source_nodes: | |
| # todo how to go from node id to document name? | |
| print(f"Fetched node {node} in search.") | |
| response_text = [] | |
| for text in response.response_gen: | |
| response_text.append(text) | |
| string_output = ''.join(response_text) | |
| yield string_output | |
| logger.info(f'Assistant response: {string_output}') | |