| import logging | |
| import os | |
| from llama_index.core import SimpleDirectoryReader, VectorStoreIndex | |
| from llama_index.core.chat_engine.types import AgentChatResponse | |
| from llama_index.llms.openai import OpenAI | |
| from llama_index.core import StorageContext, load_index_from_storage | |
| logger = logging.getLogger("app.llama_index") | |
| class Assistant: | |
| def __init__( | |
| self, | |
| data_dir: str = "data/pdf", | |
| model_name: str = "gpt-4", | |
| ) -> None: | |
| self.data_dir = data_dir | |
| self.model_name = model_name | |
| storage_dir = os.path.join(self.data_dir, "storage") | |
| if os.path.isdir(storage_dir): | |
| logger.info("Loading index from storage") | |
| storage_context = StorageContext.from_defaults(persist_dir=storage_dir) | |
| self.index = load_index_from_storage(storage_context) | |
| else: | |
| self.documents = SimpleDirectoryReader(self.data_dir).load_data() | |
| self.index = VectorStoreIndex(self.documents) | |
| self.index.storage_context.persist(storage_dir) | |
| llm = OpenAI(model=self.model_name) | |
| self.chat_engine = self.index.as_chat_engine( | |
| chat_mode="openai", | |
| llm=llm, | |
| verbose=True, | |
| system_prompt=( | |
| "Du er en assistent til at hjælpe med at besvare spørgsmål omkring årsopgørelsen fra skat." | |
| " Din viden er baseret på årsopgørelsen og relaterede emner fra skat.dk." | |
| " Du kan svare på spørgsmål omkring årsopgørelsen, og du kan også stille spørgsmål til brugeren for at få mere information." | |
| " Du kan også bede brugeren om at uddybe, hvis du ikke forstår spørgsmålet." | |
| " Du skal opgive de kilder, du bruger til at besvare spørgsmålet, og du skal også oplyse, hvis du ikke kan besvare spørgsmålet." | |
| ), | |
| ) | |
| def __call__(self, question: str, history: list[str]) -> str: | |
| response: AgentChatResponse = self.chat_engine.chat( | |
| question, | |
| tool_choice="query_engine_tool", | |
| ) | |
| return response.response | |