feat: Langchain HN RAG demo

indexing.py

@@ -0,0 +1,114 @@
+import os
+from typing import TypedDict
+
+import langchain
+import langchain.chat_models
+import langchain.hub
+import langchain.text_splitter
+import langchain_core
+import langchain_core.documents
+import langchain_core.vectorstores
+import langchain_openai
+import langgraph
+import langgraph.graph
+import mlflow
+
+from hn import HackerNewsClient, Story
+from scrape import JinaScraper
+
+
+async def fetch_hn_top_stories(
+    limit: int = 10,
+) -> list[langchain_core.documents.Document]:
+    hn = HackerNewsClient()
+    stories = hn.get_top_stories(limit=limit)
+
+    contents = {}
+
+    # Fetch content for each story asynchronously
+    scraper = JinaScraper(os.getenv("JINA_API_KEY"))
+
+    async def _fetch_content(story: Story) -> tuple[str, str]:
+        if not story.url:
+            return story.id, story.title
+        return story.id, await scraper.get_content(story.url)
+
+    tasks = [_fetch_content(story) for story in stories]
+    results = await asyncio.gather(*tasks)
+    contents = dict(results)
+
+    documents = [
+        langchain_core.documents.Document(
+            page_content=contents[story.id],
+            metadata={
+                "id": story.id,
+                "title": story.title,
+                "source": story.url,
+                "created_at": story.created_at.isoformat(),
+            },
+        )
+        for story in stories
+    ]
+    return documents
+
+
+async def main():
+    mlflow.set_tracking_uri("http://localhost:5000")
+    mlflow.set_experiment("langchain-rag-hn")
+    mlflow.langchain.autolog()
+
+    llm = langchain.chat_models.init_chat_model(
+        model="gpt-4o-mini", model_provider="openai"
+    )
+    embeddings = langchain_openai.OpenAIEmbeddings(model="text-embedding-3-small")
+    vector_store = langchain_core.vectorstores.InMemoryVectorStore(embeddings)
+
+    # 1. Load
+    stories = await fetch_hn_top_stories(limit=20)
+
+    # 2. Split
+    splitter = langchain.text_splitter.RecursiveCharacterTextSplitter(
+        chunk_size=1000, chunk_overlap=200
+    )
+    all_splits = splitter.split_documents(stories)
+
+    # 3. Store
+    _ = vector_store.add_documents(all_splits)
+
+    # 4. Query
+    prompt = langchain.hub.pull("rlm/rag-prompt")
+
+    # Define state for application
+    class State(TypedDict):
+        question: str
+        context: list[langchain_core.documents.Document]
+        answer: str
+
+    # Define application steps
+    def retrieve(state: State):
+        retrieved_docs = vector_store.similarity_search(state["question"], k=10)
+        return {"context": retrieved_docs}
+
+    def generate(state: State):
+        docs_content = "\n\n".join(doc.page_content for doc in state["context"])
+        messages = prompt.invoke(
+            {"question": state["question"], "context": docs_content}
+        )
+        response = llm.invoke(messages)
+        return {"answer": response.content}
+
+    # Compile application and test
+    graph_builder = langgraph.graph.StateGraph(State).add_sequence([retrieve, generate])
+    graph_builder.add_edge(langgraph.graph.START, "retrieve")
+    graph = graph_builder.compile()
+
+    response = graph.invoke(
+        {"question": "Are there any news stories related to AI and Machine Learning?"}
+    )
+    print(response["answer"])
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())