ragaai-catalyst 2.1.5b15__tar.gz → 2.1.5b16__tar.gz

{ragaai_catalyst-2.1.5b15 → ragaai_catalyst-2.1.5b16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ragaai_catalyst
-Version: 2.1.5b15
+Version: 2.1.5b16
 Summary: RAGA AI CATALYST
 Author-email: Kiran Scaria <kiran.scaria@raga.ai>, Kedar Gaikwad <kedar.gaikwad@raga.ai>, Dushyant Mahajan <dushyant.mahajan@raga.ai>, Siddhartha Kosti <siddhartha.kosti@raga.ai>, Ritika Goel <ritika.goel@raga.ai>, Vijay Chaurasia <vijay.chaurasia@raga.ai>
 Requires-Python: <3.13,>=3.9

ragaai_catalyst-2.1.5b16/examples/langgraph_examples/agentic_rag.py

@@ -0,0 +1,318 @@
+from langchain_community.document_loaders import WebBaseLoader
+from langchain_community.vectorstores import Chroma
+from langchain_openai import OpenAIEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from dotenv import load_dotenv
+import os
+load_dotenv()
+from ragaai_catalyst.tracers import Tracer
+from ragaai_catalyst import RagaAICatalyst, init_tracing
+from ragaai_catalyst import trace_tool, current_span, trace_agent, trace_llm
+
+catalyst = RagaAICatalyst(
+    access_key=os.getenv("RAGAAI_CATALYST_ACCESS_KEY"),
+    secret_key=os.getenv("RAGAAI_CATALYST_SECRET_KEY"),
+    base_url=os.getenv("RAGAAI_CATALYST_BASE_URL"),
+)
+
+# Initialize tracer
+tracer = Tracer(
+    project_name="Langgraph_testing",
+    dataset_name="agentic_rag",
+    tracer_type="Agentic",
+)
+
+init_tracing(catalyst=catalyst, tracer=tracer)
+
+urls = [
+    "https://lilianweng.github.io/posts/2023-06-23-agent/",
+    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
+    "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
+]
+
+docs = [WebBaseLoader(url).load() for url in urls]
+docs_list = [item for sublist in docs for item in sublist]
+
+text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+    chunk_size=100, chunk_overlap=50
+)
+doc_splits = text_splitter.split_documents(docs_list)
+
+# Add to vectorDB
+vectorstore = Chroma.from_documents(
+    documents=doc_splits,
+    collection_name="rag-chroma",
+    embedding=OpenAIEmbeddings(),
+)
+retriever = vectorstore.as_retriever()
+
+from langchain.tools.retriever import create_retriever_tool
+
+retriever_tool = create_retriever_tool(
+    retriever,
+    "retrieve_blog_posts",
+    "Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.",
+)
+
+tools = [retriever_tool]
+
+from typing import Annotated, Sequence
+from typing_extensions import TypedDict
+
+from langchain_core.messages import BaseMessage
+
+from langgraph.graph.message import add_messages
+
+
+class AgentState(TypedDict):
+    # The add_messages function defines how an update should be processed
+    # Default is to replace. add_messages says "append"
+    messages: Annotated[Sequence[BaseMessage], add_messages]
+    
+from typing import Annotated, Literal, Sequence
+from typing_extensions import TypedDict
+
+from langchain import hub
+from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import PromptTemplate
+from langchain_openai import ChatOpenAI
+
+from pydantic import BaseModel, Field
+
+
+from langgraph.prebuilt import tools_condition
+
+### Edges
+
+@trace_agent("grade_documents")
+def grade_documents(state) -> Literal["generate", "rewrite"]:
+    """
+    Determines whether the retrieved documents are relevant to the question.
+
+    Args:
+        state (messages): The current state
+
+    Returns:
+        str: A decision for whether the documents are relevant or not
+    """
+
+    print("---CHECK RELEVANCE---")
+
+    # Data model
+    class grade(BaseModel):
+        """Binary score for relevance check."""
+
+        binary_score: str = Field(description="Relevance score 'yes' or 'no'")
+
+    # LLM
+    model = ChatOpenAI(temperature=0, model="gpt-4o", streaming=True)
+
+    # LLM with tool and validation
+    llm_with_tool = model.with_structured_output(grade)
+
+    # Prompt
+    prompt = PromptTemplate(
+        template="""You are a grader assessing relevance of a retrieved document to a user question. \n 
+        Here is the retrieved document: \n\n {context} \n\n
+        Here is the user question: {question} \n
+        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
+        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.""",
+        input_variables=["context", "question"],
+    )
+
+    # Chain
+    chain = prompt | llm_with_tool
+    current_span().add_metrics(
+            name="grade_documents",
+            score=0.9,
+            reasoning="Determine relevance of retrieved documents",
+            cost=0.01,
+            latency=0.5,
+        )
+
+    messages = state["messages"]
+    last_message = messages[-1]
+
+    question = messages[0].content
+    docs = last_message.content
+
+    scored_result = chain.invoke({"question": question, "context": docs})
+
+    score = scored_result.binary_score
+
+    if score == "yes":
+        print("---DECISION: DOCS RELEVANT---")
+        return "generate"
+
+    else:
+        print("---DECISION: DOCS NOT RELEVANT---")
+        print(score)
+        return "rewrite"
+
+
+### Nodes
+
+@tracer.trace_agent("agent")
+def agent(state):
+    """
+    Invokes the agent model to generate a response based on the current state. Given
+    the question, it will decide to retrieve using the retriever tool, or simply end.
+
+    Args:
+        state (messages): The current state
+
+    Returns:
+        dict: The updated state with the agent response appended to messages
+    """
+    print("---CALL AGENT---")
+    messages = state["messages"]
+    model = ChatOpenAI(temperature=0, streaming=True, model="gpt-4-turbo")
+    model = model.bind_tools(tools)
+    response = model.invoke(messages)
+    current_span().add_metrics(
+        name="agent",
+        score=0.9,
+        reasoning="Generating response based on current state",
+        cost=0.01,
+        latency=0.5,
+    )
+    # We return a list, because this will get added to the existing list
+    return {"messages": [response]}
+
+@trace_agent("rewrite")
+def rewrite(state):
+    """
+    Transform the query to produce a better question.
+
+    Args:
+        state (messages): The current state
+
+    Returns:
+        dict: The updated state with re-phrased question
+    """
+
+    print("---TRANSFORM QUERY---")
+    messages = state["messages"]
+    question = messages[0].content
+
+    msg = [
+        HumanMessage(
+            content=f""" \n 
+    Look at the input and try to reason about the underlying semantic intent / meaning. \n 
+    Here is the initial question:
+    \n ------- \n
+    {question} 
+    \n ------- \n
+    Formulate an improved question: """,
+        )
+    ]
+
+    # Grader
+    model = ChatOpenAI(temperature=0, model="gpt-4-0125-preview", streaming=True)
+    response = model.invoke(msg)
+    return {"messages": [response]}
+
+@trace_tool("generate")
+def generate(state):
+    """
+    Generate answer
+
+    Args:
+        state (messages): The current state
+
+    Returns:
+         dict: The updated state with re-phrased question
+    """
+    print("---GENERATE---")
+    messages = state["messages"]
+    question = messages[0].content
+    last_message = messages[-1]
+
+    docs = last_message.content
+
+    # Prompt
+    prompt = hub.pull("rlm/rag-prompt")
+
+    # LLM
+    llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0, streaming=True)
+
+    # Post-processing
+    def format_docs(docs):
+        return "\n\n".join(doc.page_content for doc in docs)
+
+    # Chain
+    rag_chain = prompt | llm | StrOutputParser()
+    current_span().add_metrics(
+            name="generate",
+            score=0.9,
+            reasoning="Generating response based on current state",
+            cost=0.01,
+            latency=0.5,
+        )
+
+    # Run
+    response = rag_chain.invoke({"context": docs, "question": question})
+    return {"messages": [response]}
+
+
+print("*" * 20 + "Prompt[rlm/rag-prompt]" + "*" * 20)
+prompt = hub.pull("rlm/rag-prompt").pretty_print()  # Show what the prompt looks like
+
+from langgraph.graph import END, StateGraph, START
+from langgraph.prebuilt import ToolNode
+
+# Define a new graph
+workflow = StateGraph(AgentState)
+
+# Define the nodes we will cycle between
+workflow.add_node("agent", agent)  # agent
+retrieve = ToolNode([retriever_tool])
+workflow.add_node("retrieve", retrieve)  # retrieval
+workflow.add_node("rewrite", rewrite)  # Re-writing the question
+workflow.add_node(
+    "generate", generate
+)  # Generating a response after we know the documents are relevant
+# Call agent node to decide to retrieve or not
+workflow.add_edge(START, "agent")
+
+# Decide whether to retrieve
+workflow.add_conditional_edges(
+    "agent",
+    # Assess agent decision
+    tools_condition,
+    {
+        # Translate the condition outputs to nodes in our graph
+        "tools": "retrieve",
+        END: END,
+    },
+)
+
+# Edges taken after the `action` node is called.
+workflow.add_conditional_edges(
+    "retrieve",
+    # Assess agent decision
+    grade_documents,
+)
+workflow.add_edge("generate", END)
+workflow.add_edge("rewrite", "agent")
+
+# Compile
+graph = workflow.compile()
+
+import pprint
+
+inputs = {
+    "messages": [
+        ("user", "What did Lilian Weng write about automatic prompt generation?"),
+        ("user", "What did Lilian Weng write about attack on LLMs?"),
+    ]
+}
+
+with tracer:
+    for output in graph.stream(inputs):
+        for key, value in output.items():
+            pprint.pprint(f"Output from node '{key}':")
+            pprint.pprint("---")
+            pprint.pprint(value, indent=2, width=80, depth=None)
+        pprint.pprint("\n---\n")