[feat] add find files methods

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
+<pre><code>```mermaid
+graph TD
+	__start__([<p>__start__</p>]):::first
+	find_references(find_references)
+	agent(agent)
+	tools(tools)
+	crag(crag)
+	generate(generate)
+	__end__([<p>__end__</p>]):::last
+	crag -->generate;
+	find_references --> generate;
+	generate --> __end__;
+	tools --> crag;
+	__start__ -. &nbsp;general&nbsp; .-> agent;
+	__start__ -. &nbsp;specific&nbsp; .-> find_references;
+	agent -.&nbsp;continue&nbsp; .-> tools;
+	agent -. &nbsp;end&nbsp; .-> __end__
+```</code></pre>
+
+
+
 ## 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/)
+

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(

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,
         )
 

src/infrastructure/database/chromadb/connector.py

@@ -143,21 +143,48 @@ 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
         """
         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}"
+            )

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,30 +33,49 @@ 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}")
 
     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()