diff --git a/README.md b/README.md index d0e05da..c3f131c 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# CRAG API +# Bot de Notícias ![Python](https://img.shields.io/badge/Python-3776AB?style=for-the-badge&logo=python&logoColor=white) ![Docker](https://img.shields.io/badge/Docker-2496ED?style=for-the-badge&logo=docker&logoColor=white) @@ -9,10 +9,11 @@ ![ChromaDB](https://img.shields.io/badge/ChromaDB-FFA500?style=for-the-badge&logo=prisma&logoColor=white) ![Llama](https://img.shields.io/badge/Llama-FF6B6B?style=for-the-badge&logo=meta&logoColor=white) ![OpenAI](https://img.shields.io/badge/OpenAI-412991?style=for-the-badge&logo=openai&logoColor=white) +![C++](https://img.shields.io/badge/C++-00599C?style=for-the-badge&logo=c%2B%2B&logoColor=white) ## Sobre o Projeto -A CRAG API é uma aplicação conversacional baseada em um grafo que utiliza a técnica de Retrieval-Augmented Generation (RAG), com um nó adicional de correção. Antes de gerar uma resposta com base em um documento semanticamente semelhante, o sistema avalia seu contexto, garantindo maior precisão e coerência. +É aplicação conversacional baseada em um grafo que utiliza a técnica de Retrieval-Augmented Generation (RAG), com um nó adicional de correção. Antes de gerar uma resposta com base em um documento semanticamente semelhante, o sistema avalia seu contexto, garantindo maior precisão e coerência. O sistema adota uma arquitetura em camadas (Layered Architecture), utilizando FastAPI como framework principal para a API REST. O armazenamento é dividido entre ChromaDB, responsável pela gestão dos arquivos usados no RAG, e MongoDB, que armazena os logs da aplicação. @@ -26,6 +27,28 @@ docker-compose -f docker/docker-compose.yml --env-file .env up --build ``` PS: Não se esqueça de alterar os [Prompts](src/services/crag/prompts.py). +## O grafo +
```mermaid
+graph TD
+	__start__([
__start__
]):::first
+	find_references(find_references)
+	agent(agent)
+	tools(tools)
+	crag(crag)
+	generate(generate)
+	__end__([
__end__
]):::last
+	crag -->generate;
+	find_references --> generate;
+	generate --> __end__;
+	tools --> crag;
+	__start__ -.  general  .-> agent;
+	__start__ -.  specific  .-> find_references;
+	agent -. continue  .-> tools;
+	agent -.  end  .-> __end__
+```
+ + + ## Autor **[@CuriousGu](https://www.github.com/CuriousGu) 🇧🇷** @@ -46,3 +69,4 @@ Veja o arquivo [LICENSE](LICENSE) para mais detalhes. ## Contatos - Email: gustavo_ortega@usp.br - Linkedin: [Gustavo M. Ortega](https://www.linkedin.com/in/gustavomendoncaortega/) + diff --git a/src/api/controllers/crag.py b/src/api/controllers/crag.py index 068ceb7..e53865a 100644 --- a/src/api/controllers/crag.py +++ b/src/api/controllers/crag.py @@ -33,7 +33,9 @@ async def contr_new_message( history.append({ "role": "assistant", "content": response["messages"], - "timestamp": datetime.now().isoformat() + "timestamp": datetime.now().isoformat(), + "docs": response.get("docs", []), + "decision_type": response.get("decision_type", "") }) _ = await add_message_to_history( diff --git a/src/api/routes/files.py b/src/api/routes/files.py index 2cec925..d560a1c 100644 --- a/src/api/routes/files.py +++ b/src/api/routes/files.py @@ -82,7 +82,7 @@ async def upload_files( try: _ = await controller_upload_file( file=file, - metadata=FileMetadata(), + metadata=FileMetadata(file_name=file.filename), vector_store=req.app.vector_store, ) diff --git a/src/infrastructure/database/chromadb/connector.py b/src/infrastructure/database/chromadb/connector.py index 51e2a30..253dbd4 100644 --- a/src/infrastructure/database/chromadb/connector.py +++ b/src/infrastructure/database/chromadb/connector.py @@ -143,13 +143,14 @@ async def close(self): @staticmethod @tool("retriever") - def retrieve(query: str) -> List: + def retrieve(query: str, n: int = 10) -> List: """ Método que faz uma requisição a vector store para consultar os documentos que podem ajudar a responder a pergunta do usuário. Args: query (str): A query para a vector store. + n (int): Quantidade de documentos a serem retornados. Returns: List[Document]: Uma lista de documentos recuperados @@ -157,7 +158,33 @@ def retrieve(query: str) -> List: try: db = ChromaDB() retriever = db._as_retriever(settings.INDEX_NAME) - return retriever.invoke(query) + # Limitando a quantidade de documentos retornados para n+10 + # Isso garante que tenhamos documentos suficientes para filtrar posteriormente + response = retriever.invoke(query) + + created_at = [ + (index, x["created_at"]) + for index, x in enumerate(response.get("metadatas", [])) + ] + created_at.sort(key=lambda x: x[1]) + + unique_titles = set() + unique_documents = [] + + for index, _ in created_at: + title = response["metadatas"][index].get("file_name") + if title and title not in unique_titles: + unique_titles.add(title) + unique_documents.append( + { + "page_content": response["documents"][index], + "metadata": {"file_name": title} + } + ) + if len(unique_documents) == n: + break + + return unique_documents except Exception as e: raise e @@ -166,7 +193,7 @@ def retrieve(query: str) -> List: def get_most_recent(n: int = 5) -> List: """ Método que faz uma requisição a vector store para consultar os - arquivos que foram adicionados mais recentemente. + arquivos SEM CONTEXTO ESPECÍFICO, que foram adicionados recentemente. Args: n (int): Quantidade de arquivos a serem retornados. @@ -185,10 +212,49 @@ def get_most_recent(n: int = 5) -> List: ] created_at.sort(key=lambda x: x[1]) - return [ - {"page_content": results["documents"][index]} - for index, _ in created_at[-n:] - ] + unique_titles = set() + unique_documents = [] + + for index, _ in created_at: + title = results["metadatas"][index].get("file_name") + if title and title not in unique_titles: + unique_titles.add(title) + unique_documents.append( + { + "page_content": results["documents"][index], + "metadata": {"file_name": title} + } + ) + if len(unique_documents) == n: + break + + return unique_documents except Exception as e: raise ValueError(f"Erro ao buscar arquivos recentes: {e}") + + @staticmethod + def find_files(file_name: str) -> List: + try: + db = ChromaDB() + collection = db.client.get_collection(settings.INDEX_NAME) + results = collection.get() or [] + + found_documents = [] + + for index, metadata in enumerate(results.get("metadatas", [])): + title = metadata.get("file_name", "") + if title and file_name.lower() in title.lower(): + found_documents.append( + { + "page_content": results["documents"][index], + "metadata": metadata + } + ) + + return found_documents + + except Exception as e: + raise ValueError( + f"Erro ao buscar arquivos por nome: {e}" + ) diff --git a/src/services/crag/graph.py b/src/services/crag/graph.py index 265de93..ccd9070 100644 --- a/src/services/crag/graph.py +++ b/src/services/crag/graph.py @@ -6,10 +6,12 @@ from src.infrastructure.config import settings from .templates import AgentState from .nodes import ( - agent, + get_news, should_continue, generate, grade_documents, + find_references, + flow_decision, CustomToolNode ) @@ -31,7 +33,11 @@ async def invoke( "model": model, } ) - return {"messages": response["messages"][-1].content} + return { + "messages": response["messages"][-1].content, + "docs": list(response["docs"]) if response.get("docs") else [], + "decision_type": response["decision_type"] + } except Exception as e: raise ValueError(f"Error invoking CRAG: {e}") @@ -39,22 +45,37 @@ async def invoke( def build(self): try: builder = StateGraph(AgentState) - builder.add_node("agent", lambda state: agent(state)) + builder.add_node("find_references", find_references) + builder.add_node("get_news", get_news) builder.add_node("tools", CustomToolNode()) builder.add_node("crag", grade_documents) builder.add_node("generate", generate) - builder.add_edge(START, "agent") builder.add_conditional_edges( - "agent", + START, + flow_decision, + { + "general": "get_news", + "specific": "find_references", + "end": "generate" + } + ) + # DETAILED SEARCH + builder.add_edge("find_references", "generate") + + # MOST RECENT OR RETRIEVER + builder.add_conditional_edges( + "get_news", should_continue, { "continue": "tools", - "end": END + "end": "generate" } ) builder.add_edge("tools", "crag") builder.add_edge("crag", "generate") + + # END builder.add_edge("generate", END) self.graph = builder.compile() diff --git a/src/services/crag/nodes.py b/src/services/crag/nodes.py index 80d7203..02de667 100644 --- a/src/services/crag/nodes.py +++ b/src/services/crag/nodes.py @@ -1,51 +1,128 @@ - from langchain_core.messages import ToolMessage, SystemMessage from langchain_core.prompts import PromptTemplate +import html -from .templates import AgentState, GradeDocument +from .templates import AgentState, GradeDocument, FlowDecision, IsSpecificFile from src.infrastructure.database import ChromaDB from .prompts import ( - grader_prompt, agent_prompt, no_generation, generate_answer_prompt + grader_prompt, + get_news_prompt, + no_generation, + generate_answer_prompt, + flow_decision_prompt, + is_specific_file ) -class CustomToolNode: - def __init__(self): - self.db = ChromaDB() - self.tools = { - "retriever": ChromaDB().retrieve, - "most_recent_files": ChromaDB().get_most_recent +def flow_decision(state: AgentState): + # decide se vai recuperar as informações de + # um modo geral ou procurar informações mais específicas + # sobre uma noticia específica. + LLM = state["model"] + flow_decision = PromptTemplate.from_template( + flow_decision_prompt + ) + chain = ( + flow_decision | + LLM.with_structured_output(FlowDecision) + ) + + formatted_history = "\n".join( + [ + f"{'user' if msg.type == 'human' else 'IA'}: " \ + f"{msg.content}" for msg in state["messages"][-5:] + ] + ) + + response = chain.invoke( + { + "history": formatted_history, + "question": state["messages"][-1].content } + ) - def __call__(self, inputs: list): - if messages := inputs.get("messages", []): - message = messages[-1] - else: - raise ValueError("No messages found in inputs") + return response.decision if hasattr(response, "decision") else "end" - for tool_call in message.tool_calls: - tool_result = self.tools[tool_call["name"]].invoke( - tool_call["args"] - ) - if not isinstance(tool_call["args"], dict): - raise TypeError("Tool call args must be a dictionary") +def find_references(state: AgentState): + # puxa a ultima mensagem do tipo generic + # pega as referencias das noticias + # consulta qual delas tem relacao com a pergunta do user. + # submete os arquivos relevantes para o generate + messages = state["messages"] + LLM = state["model"] - query = tool_call["args"].get("query") - docs = tool_result + is_specific_file_chain = ( + PromptTemplate.from_template(is_specific_file) + | LLM.with_structured_output(IsSpecificFile) + ) - messages.append( - ToolMessage( - content="Documentos encontrados e adicionardos ao contexto", - tool_call_id=tool_call["id"] - ) + for i in messages[::-1]: + if ( + i.additional_kwargs.get("docs") and + i.additional_kwargs.get("decision_type") == "general" + ): + docs_last_message = i.additional_kwargs.get("docs") + break + + formatted_history = "\n".join( + [ + f"{'user' if msg.type == 'human' else 'IA'}: " \ + f"{msg.content}" for msg in state["messages"][-5:] + ] + ) + + # consulta ao documento + for doc in docs_last_message: + found_docs = ChromaDB().find_files(doc) + decoded_docs = [ + html.unescape(x["page_content"]) + .replace(" ", " ") + .replace("\\r\\n", "\\n") + .strip() + for x in found_docs + ] + + relevant = is_specific_file_chain.invoke( + { + "document": "\n".join(decoded_docs), + "question": messages[-1].content, + "history": formatted_history + } ) - return { - "query": query, - "docs": docs, - "messages": messages - } + if relevant and relevant.arquivo_relevante: + return { + "docs": found_docs, + "decision_type": "specific", + "user_last_message": messages[-1] + } + + return { + "docs": [], + "decision_type": "specific", + "user_last_message": messages[-1] + } + + +def get_news(state: AgentState): + LLM = state["model"] + agent_msg = PromptTemplate.from_template(get_news_prompt) + chain = agent_msg | LLM.bind_tools( + [ChromaDB().retrieve, ChromaDB().get_most_recent] + ) + + return { + "user_last_message": state["messages"][-1], + "messages": [chain.invoke(input={"question": state["messages"][-5:]})], + "decision_type": "general" + } + + +def should_continue(state: AgentState): + if state["messages"][-1].tool_calls: + return "continue" + return "end" def grade_documents(state: AgentState): @@ -69,9 +146,12 @@ def grade_documents(state: AgentState): if not isinstance(d, dict) else d["page_content"] ), - "message": messages} + "message": messages[:-3] + } ) - if score and score.binary_score == "yes": + if (score and score.binary_score == "yes") or ( + messages[-2].tool_calls[0]["name"] == "most_recent_files" + ): filtered_docs.append(d) return { @@ -81,31 +161,13 @@ def grade_documents(state: AgentState): } -def agent(state: AgentState): - LLM = state["model"] - agent_msg = PromptTemplate.from_template(agent_prompt) - chain = agent_msg | LLM.bind_tools( - [ChromaDB().retrieve, ChromaDB().get_most_recent] - ) - - return { - "messages": [chain.invoke(input={"question": state["messages"][-1]})] - } - - -def should_continue(state: AgentState): - if state["messages"][-1].tool_calls: - return "continue" - return "end" - - def generate(state: AgentState): LLM = state["model"] docs = state.get("docs", None) + user_last_message = state.get("user_last_message", None) messages = state.get("messages", []) - query = state.get("query", None) or messages[0].content - if len(docs) >= 1 and isinstance(messages[-1], ToolMessage): + if docs and len(docs) >= 1: answer_chain = ( PromptTemplate( input_variables=["query", "context", "message"], @@ -116,18 +178,66 @@ def generate(state: AgentState): result = answer_chain.invoke( { - "query": query, + "query": user_last_message, "context": docs, - "message": messages + "messages": messages } ) - return { - "messages": result, - } + response = { + "messages": result, + "docs": [ + x.metadata["file_name"] + if hasattr(x, 'metadata') + else x["metadata"]["file_name"] + for x in docs + ], + "decision_type": state["decision_type"] + } + response["docs"] = set(response["docs"]) + return response else: messages = [SystemMessage(content=no_generation)] + state["messages"] return { "messages": [LLM.invoke(messages)], + "decision_type": "no_generation" + } + + +class CustomToolNode: + def __init__(self): + self.tools = { + "retriever": ChromaDB().retrieve, + "most_recent_files": ChromaDB().get_most_recent + } + + def __call__(self, inputs: list): + if messages := inputs.get("messages", []): + message = messages[-1] + else: + raise ValueError("No messages found in inputs") + + for tool_call in message.tool_calls: + tool_result = self.tools[tool_call["name"]].invoke( + tool_call["args"] + ) + + if not isinstance(tool_call["args"], dict): + raise TypeError("Tool call args must be a dictionary") + + query = tool_call["args"].get("query") + docs = tool_result + + messages.append( + ToolMessage( + content="Documentos encontrados e adicionardos ao contexto", + tool_call_id=tool_call["id"] + ) + ) + + return { + "query": query, + "docs": docs, + "messages": messages } diff --git a/src/services/crag/prompts.py b/src/services/crag/prompts.py index 11772e2..eaa56a2 100644 --- a/src/services/crag/prompts.py +++ b/src/services/crag/prompts.py @@ -1,43 +1,159 @@ -agent_prompt = """ - Gere uma consulta para a vector store com base na pergunta do usuario: - {question} +flow_decision_prompt = """ + Você é um assistente inteligente que classifica perguntas de usuários em três categorias: "specific", "general" e "end". + + 1. **"specific"**: Use esta categoria quando a pergunta do usuário se refere a uma notícia específica ou a um evento que foi mencionado anteriormente na conversa. + 2. **"general"**: Use esta categoria quando a pergunta do usuário é sobre notícias em geral, mas não se refere a uma notícia específica ou a um evento já discutido. + 3. **"end"**: Use esta categoria quando a pergunta do usuário não está relacionada a notícias ou eventos. + + Considere as interações anteriores para determinar se a pergunta é sobre uma notícia específica ou não. Responda apenas com a categoria correspondente. + + **Exemplos de interações:** + + exemplo 1: + - Usuário: "Quais as notícias de hoje?" + Resposta: "general" + + exemplo 2: + - Usuário: "Competição de xadrez" + Resposta: "general" + + exemplo 3: + - Usuário: "oi" + Resposta: "end" + + exemplo 4: + - Usuário: "O que é o xadrez?" + Resposta: "end" + + exemplo 5: + - Usuário: "Quais as notícias de hoje?" + - AImessage: "Campeonato de xadrez na USP dia 20/04" + - Usuário: "Me fale sobre a competição de xadrez" + Resposta: "specific" + + exemplo 6: + - Usuário: "Quais as notícias de hoje?" + - AImessage: "Vai acontecer uma competição de robótica na USP" + - Usuário: "Me fale sobre a competição de robótica?" + Resposta: "specific" + + **Agora, classifique a seguinte pergunta do usuário:** + {question} - Para ajudar a responder a pergunta, considere as seguintes ferramentas: - - Most Recent: Caso seja sobre os documentos mais recentes, use - essa ferramenta para obter os documentos mais recentes. - - Retriever: Consulta a vector store para buscar contexto - sobre a questão, exceto se a pergunta for sobre os documentos - mais recentes. + **Histórico de mensagens:** + {history} """ -grader_prompt = """ - Avalie os arquivos recuperados com base na pergunta do usuario, - sendo que sua resposta deve ser "yes", caso relevante, ou "no": +is_specific_file = """ + Você é um avaliador especializado em determinar a relevância de um texto em relação a uma pergunta específica do usuário. Sua tarefa é analisar o documento fornecido e decidir se ele responde claramente à pergunta do usuário. + + **Critérios de Avaliação**: + - Responda **True** se o texto aborda diretamente a pergunta do usuário e fornece informações relevantes. + - Responda **False** se o texto não se relaciona ou não responde à pergunta do usuário de forma clara. + - Considere o histórico de mensagens para entender o contexto da pergunta do usuário e a intenção por trás dela. + + **Instruções**: + 1. Leia atentamente a pergunta do usuário e o texto a ser avaliado. + 2. Avalie se o texto contém informações que respondem diretamente à pergunta. + 3. Se necessário, considere o histórico de mensagens para captar nuances ou intenções que possam influenciar a relevância. + + **Pergunta do Usuário**: {question} + + **Texto a ser Avaliado**: {document} - {message} + + **Histórico de Mensagens**: + {history} + + **Por favor, forneça sua resposta apenas como "True" ou "False".** +""" + + +get_news_prompt = """ + Você é um orquestrador de ferramentas, responsável por identificar a ferramenta mais adequada a ser chamada com base na pergunta do usuário. Sua tarefa é analisar a pergunta e decidir qual ferramenta utilizar: + + **Pergunta do Usuário:** + {question} + + **Ferramentas Disponíveis:** + + 1. **Retriever**: Para pedidos de notícias com contexto específico. + 2. **Most Recent**: Para pedidos de notícias sem contexto específico. + + **Instruções:** + + - Analise a pergunta do usuário cuidadosamente. + - Retorne a chamada da ferramenta (tool call). + + **Exemplos:** + + exemplo 1: + - Pergunta: "Quais as últimas notícias sobre Programação?" + - Resposta: "Retriever" + - Justificativa: Mesmo que trata-se das notícias mais recentes, é pedido um tema específico, portanto, a ferramenta "Retriever" é a mais adequada. + + exemplo 2: + - Pergunta: "Quais as últimas notícias?" + - Resposta: "Most Recent" + - Justificativa: A pergunta não especifica um tema específico, portanto, a ferramenta "Most Recent" é a mais adequada. + + exemplo 3: + - Pergunta: "Quero notícias recentes sobre o xadrez" + - Resposta: "Retriever" + - Justificativa: Mesmo que trata-se das notícias mais recentes, é pedido um tema específico, portanto, a ferramenta "Retriever" é a mais adequada. + + exemplo 4: + - Pergunta: "Notícias sobre Hackathon" + - Resposta: "Retriever" + - Justificativa: A pergunta especifica um tema específico, portanto, a ferramenta "Retriever" é a mais adequada. + +""" + + +grader_prompt = """ + Você é um avaliador que deve avaliar se o documento recuperado é + relevante para a pergunta do usuario. + + ## Avaliação + - Se o documento for relevante para a pergunta do usuario, sua + resposta deve ser "yes". + - Caso contrário, sua resposta deve ser "no". + + **pergunta**: + {question} + **documento**: + {document} + **histórico de mensagens**: + {message} OBS: - - Se houver uma tool call para a ferramenta "Most Recent", - a resposta deve ser OBRIGATORIAMENTE "yes". + - Se a última mensagem do histórico de mensagens for uma + tool call para a ferramenta "Most Recent", a resposta deve ser + OBRIGATORIAMENTE "yes". """ no_generation = """ - Informe ao usuario que não foi possivel gerar uma resposta para - a sua pergunta. + # TAREFA + Responda de maneira objetiva ao usuário. + + # OBS + Caso não tenha contexto suficiente para responder, responda + que não foi possivel gerar uma resposta para a sua pergunta. """ generate_answer_prompt = """ - Escreva uma resposta para a pergunta do usuario com base nos - arquivos recuperados: - {query} - {context} - {message} + Gere uma resposta clara e concisa para a pergunta do usuário, utilizando as informações dos documentos recuperados. A resposta deve ser formatada para fácil leitura no Telegram e deve seguir o template abaixo: - OBS: - - A resposta deve ser gerada na mesma lingua da entrada abaixo: - {query} + **Pergunta do Usuário**: {query} + **Histórico de Mensagens**: {messages} + **Contexto**: {context} + + **Instruções**: + - A resposta deve ser em PORTUGUÊS. + - forneça uma resposta direta e informativa, mantendo a clareza e a objetividade. + - Caso a resposta seja grande, quebre em bullets. """ diff --git a/src/services/crag/templates.py b/src/services/crag/templates.py index a0c5b34..088d959 100644 --- a/src/services/crag/templates.py +++ b/src/services/crag/templates.py @@ -15,14 +15,27 @@ class AgentState(TypedDict): docs: List[Dict] model: OllamaLLM | ChatOpenAI index_name: str = Field(default=settings.INDEX_NAME) + user_last_message: str + decision_type: str class GradeDocument(BaseModel): - """ Binary Score to evaluate the relevance of the text to the query + """ Evaluate the relevance of the text to the query The response values must be "yes" or "no". These documents must support the agent's response. """ binary_score: str = Field( ..., - description="binary score of the document relevance" + description=""" + binary score of the document relevance + The response values must be "yes" or "no". + """ ) + + +class FlowDecision(BaseModel): + decision: str + + +class IsSpecificFile(BaseModel): + arquivo_relevante: bool diff --git a/src/services/document_reader/reader.py b/src/services/document_reader/reader.py index 43cb684..19cd2dd 100644 --- a/src/services/document_reader/reader.py +++ b/src/services/document_reader/reader.py @@ -53,5 +53,5 @@ async def _read_pdf(contents): return pdf_text @staticmethod - async def _read_plain(contents): + async def _read_txt(contents): return contents.decode('utf-8')