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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (104) hide show
  1. {ragaai_catalyst-2.1.5b14/ragaai_catalyst.egg-info → ragaai_catalyst-2.1.5b15}/PKG-INFO +2 -1
  2. ragaai_catalyst-2.1.5b15/examples/agentic_rag.py +320 -0
  3. ragaai_catalyst-2.1.5b15/examples/customer_support.py +556 -0
  4. ragaai_catalyst-2.1.5b15/examples/planning_agent.py +195 -0
  5. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/pyproject.toml +2 -1
  6. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py +17 -21
  7. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py +16 -8
  8. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py +52 -6
  9. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py +57 -13
  10. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py +130 -18
  11. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py +93 -0
  12. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json +35 -158
  13. ragaai_catalyst-2.1.5b15/ragaai_catalyst/tracers/agentic_tracing/utils/supported_llm_provider.toml +34 -0
  14. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/langchain_callback.py +4 -0
  15. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/tracer.py +45 -26
  16. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15/ragaai_catalyst.egg-info}/PKG-INFO +2 -1
  17. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst.egg-info/SOURCES.txt +4 -0
  18. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst.egg-info/requires.txt +1 -0
  19. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/.github/ISSUE_TEMPLATE/bug_report.md +0 -0
  20. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/.github/ISSUE_TEMPLATE/feature_request.md +0 -0
  21. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/.github/PULL_REQUEST_TEMPLATE.md +0 -0
  22. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/.gitignore +0 -0
  23. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/LICENSE +0 -0
  24. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/README.md +0 -0
  25. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/docs/dataset_management.md +0 -0
  26. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/docs/prompt_management.md +0 -0
  27. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/FinancialAnalysisSystem.ipynb +0 -0
  28. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/TravelPlanner.ipynb +0 -0
  29. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/custom_tracer_example.py +0 -0
  30. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/finance.py +0 -0
  31. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/prompt_management_litellm.ipynb +0 -0
  32. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/prompt_management_openai.ipynb +0 -0
  33. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/sync_sample_call.py +0 -0
  34. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/travel_agent/agents.py +0 -0
  35. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/travel_agent/config.py +0 -0
  36. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/travel_agent/main.py +0 -0
  37. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/examples/travel_agent/tools.py +0 -0
  38. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/__init__.py +0 -0
  39. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/_version.py +0 -0
  40. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/dataset.py +0 -0
  41. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/evaluation.py +0 -0
  42. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/experiment.py +0 -0
  43. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/guard_executor.py +0 -0
  44. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/guardrails_manager.py +0 -0
  45. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/internal_api_completion.py +0 -0
  46. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/prompt_manager.py +0 -0
  47. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/proxy_call.py +0 -0
  48. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/ragaai_catalyst.py +0 -0
  49. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/synthetic_data_generation.py +0 -0
  50. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/__init__.py +0 -0
  51. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/README.md +0 -0
  52. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/__init__.py +0 -0
  53. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/data/__init__.py +0 -0
  54. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/data/data_structure.py +0 -0
  55. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tests/FinancialAnalysisSystem.ipynb +0 -0
  56. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tests/GameActivityEventPlanner.ipynb +0 -0
  57. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tests/TravelPlanner.ipynb +0 -0
  58. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tests/__init__.py +0 -0
  59. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tests/ai_travel_agent.py +0 -0
  60. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tests/unique_decorator_test.py +0 -0
  61. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/__init__.py +0 -0
  62. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/base.py +0 -0
  63. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/langgraph_tracer.py +0 -0
  64. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/network_tracer.py +0 -0
  65. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/tracers/user_interaction_tracer.py +0 -0
  66. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/upload/__init__.py +0 -0
  67. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/upload/upload_agentic_traces.py +0 -0
  68. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/upload/upload_code.py +0 -0
  69. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py +0 -0
  70. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/__init__.py +0 -0
  71. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/api_utils.py +0 -0
  72. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py +0 -0
  73. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/file_name_tracker.py +0 -0
  74. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/generic.py +0 -0
  75. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/get_user_trace_metrics.py +0 -0
  76. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/span_attributes.py +0 -0
  77. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/system_monitor.py +0 -0
  78. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/trace_utils.py +0 -0
  79. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/unique_decorator.py +0 -0
  80. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py +0 -0
  81. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/distributed.py +0 -0
  82. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/exporters/__init__.py +0 -0
  83. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/exporters/file_span_exporter.py +0 -0
  84. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/exporters/raga_exporter.py +0 -0
  85. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/instrumentators/__init__.py +0 -0
  86. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/instrumentators/langchain.py +0 -0
  87. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/instrumentators/llamaindex.py +0 -0
  88. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/instrumentators/openai.py +0 -0
  89. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/llamaindex_callback.py +0 -0
  90. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/upload_traces.py +0 -0
  91. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/utils/__init__.py +0 -0
  92. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/utils/convert_langchain_callbacks_output.py +0 -0
  93. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/utils/langchain_tracer_extraction_logic.py +0 -0
  94. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/tracers/utils/utils.py +0 -0
  95. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst/utils.py +0 -0
  96. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst.egg-info/dependency_links.txt +0 -0
  97. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/ragaai_catalyst.egg-info/top_level.txt +0 -0
  98. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/requirements.txt +0 -0
  99. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/setup.cfg +0 -0
  100. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/test/test_catalyst/test_configuration.py +0 -0
  101. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/test/test_catalyst/test_dataset.py +0 -0
  102. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/test/test_catalyst/test_evaluation.py +0 -0
  103. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/test/test_catalyst/test_prompt_manager.py +0 -0
  104. {ragaai_catalyst-2.1.5b14 → ragaai_catalyst-2.1.5b15}/test/test_catalyst/test_synthetic_data_generation.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: ragaai_catalyst
3
- Version: 2.1.5b14
3
+ Version: 2.1.5b15
4
4
  Summary: RAGA AI CATALYST
5
5
  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>
6
6
  Requires-Python: <3.13,>=3.9
@@ -35,6 +35,7 @@ Requires-Dist: py-cpuinfo~=9.0.0
35
35
  Requires-Dist: requests~=2.32.3
36
36
  Requires-Dist: GPUtil~=1.4.0
37
37
  Requires-Dist: ipynbname
38
+ Requires-Dist: tiktoken>=0.7.0
38
39
  Provides-Extra: dev
39
40
  Requires-Dist: pytest; extra == "dev"
40
41
  Requires-Dist: pytest-cov; extra == "dev"
@@ -0,0 +1,320 @@
1
+ from langchain_community.document_loaders import WebBaseLoader
2
+ from langchain_community.vectorstores import Chroma
3
+ from langchain_openai import OpenAIEmbeddings
4
+ from langchain_text_splitters import RecursiveCharacterTextSplitter
5
+ from dotenv import load_dotenv
6
+ import os
7
+
8
+ load_dotenv()
9
+ import sys
10
+ sys.path.append("/Users/abs/Desktop/LLM/Catalyst/RagaAI-Catalyst")
11
+ from ragaai_catalyst.tracers import Tracer
12
+ from ragaai_catalyst import RagaAICatalyst, init_tracing
13
+ from ragaai_catalyst import trace_tool, current_span, trace_agent, trace_llm
14
+
15
+ catalyst = RagaAICatalyst(
16
+ access_key=os.getenv("RAGAAI_CATALYST_ACCESS_KEY"),
17
+ secret_key=os.getenv("RAGAAI_CATALYST_SECRET_KEY"),
18
+ base_url=os.getenv("RAGAAI_CATALYST_BASE_URL"),
19
+ )
20
+
21
+ # Initialize tracer
22
+ tracer = Tracer(
23
+ project_name="Trace_testing",
24
+ dataset_name="langgraph_testing",
25
+ tracer_type="Agentic",
26
+ )
27
+
28
+ init_tracing(catalyst=catalyst, tracer=tracer)
29
+
30
+ urls = [
31
+ "https://lilianweng.github.io/posts/2023-06-23-agent/",
32
+ "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
33
+ "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
34
+ ]
35
+
36
+ docs = [WebBaseLoader(url).load() for url in urls]
37
+ docs_list = [item for sublist in docs for item in sublist]
38
+
39
+ text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
40
+ chunk_size=100, chunk_overlap=50
41
+ )
42
+ doc_splits = text_splitter.split_documents(docs_list)
43
+
44
+ # Add to vectorDB
45
+ vectorstore = Chroma.from_documents(
46
+ documents=doc_splits,
47
+ collection_name="rag-chroma",
48
+ embedding=OpenAIEmbeddings(),
49
+ )
50
+ retriever = vectorstore.as_retriever()
51
+
52
+ from langchain.tools.retriever import create_retriever_tool
53
+
54
+ retriever_tool = create_retriever_tool(
55
+ retriever,
56
+ "retrieve_blog_posts",
57
+ "Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.",
58
+ )
59
+
60
+ tools = [retriever_tool]
61
+
62
+ from typing import Annotated, Sequence
63
+ from typing_extensions import TypedDict
64
+
65
+ from langchain_core.messages import BaseMessage
66
+
67
+ from langgraph.graph.message import add_messages
68
+
69
+
70
+ class AgentState(TypedDict):
71
+ # The add_messages function defines how an update should be processed
72
+ # Default is to replace. add_messages says "append"
73
+ messages: Annotated[Sequence[BaseMessage], add_messages]
74
+
75
+ from typing import Annotated, Literal, Sequence
76
+ from typing_extensions import TypedDict
77
+
78
+ from langchain import hub
79
+ from langchain_core.messages import BaseMessage, HumanMessage
80
+ from langchain_core.output_parsers import StrOutputParser
81
+ from langchain_core.prompts import PromptTemplate
82
+ from langchain_openai import ChatOpenAI
83
+
84
+ from pydantic import BaseModel, Field
85
+
86
+
87
+ from langgraph.prebuilt import tools_condition
88
+
89
+ ### Edges
90
+
91
+ @trace_agent("grade_documents")
92
+ def grade_documents(state) -> Literal["generate", "rewrite"]:
93
+ """
94
+ Determines whether the retrieved documents are relevant to the question.
95
+
96
+ Args:
97
+ state (messages): The current state
98
+
99
+ Returns:
100
+ str: A decision for whether the documents are relevant or not
101
+ """
102
+
103
+ print("---CHECK RELEVANCE---")
104
+
105
+ # Data model
106
+ class grade(BaseModel):
107
+ """Binary score for relevance check."""
108
+
109
+ binary_score: str = Field(description="Relevance score 'yes' or 'no'")
110
+
111
+ # LLM
112
+ model = ChatOpenAI(temperature=0, model="gpt-4o", streaming=True)
113
+
114
+ # LLM with tool and validation
115
+ llm_with_tool = model.with_structured_output(grade)
116
+
117
+ # Prompt
118
+ prompt = PromptTemplate(
119
+ template="""You are a grader assessing relevance of a retrieved document to a user question. \n
120
+ Here is the retrieved document: \n\n {context} \n\n
121
+ Here is the user question: {question} \n
122
+ If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
123
+ Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.""",
124
+ input_variables=["context", "question"],
125
+ )
126
+
127
+ # Chain
128
+ chain = prompt | llm_with_tool
129
+ current_span().add_metrics(
130
+ name="grade_documents",
131
+ score=0.9,
132
+ reasoning="Determine relevance of retrieved documents",
133
+ cost=0.01,
134
+ latency=0.5,
135
+ )
136
+
137
+ messages = state["messages"]
138
+ last_message = messages[-1]
139
+
140
+ question = messages[0].content
141
+ docs = last_message.content
142
+
143
+ scored_result = chain.invoke({"question": question, "context": docs})
144
+
145
+ score = scored_result.binary_score
146
+
147
+ if score == "yes":
148
+ print("---DECISION: DOCS RELEVANT---")
149
+ return "generate"
150
+
151
+ else:
152
+ print("---DECISION: DOCS NOT RELEVANT---")
153
+ print(score)
154
+ return "rewrite"
155
+
156
+
157
+ ### Nodes
158
+
159
+ @tracer.trace_agent("agent")
160
+ def agent(state):
161
+ """
162
+ Invokes the agent model to generate a response based on the current state. Given
163
+ the question, it will decide to retrieve using the retriever tool, or simply end.
164
+
165
+ Args:
166
+ state (messages): The current state
167
+
168
+ Returns:
169
+ dict: The updated state with the agent response appended to messages
170
+ """
171
+ print("---CALL AGENT---")
172
+ messages = state["messages"]
173
+ model = ChatOpenAI(temperature=0, streaming=True, model="gpt-4-turbo")
174
+ model = model.bind_tools(tools)
175
+ response = model.invoke(messages)
176
+ current_span().add_metrics(
177
+ name="agent",
178
+ score=0.9,
179
+ reasoning="Generating response based on current state",
180
+ cost=0.01,
181
+ latency=0.5,
182
+ )
183
+ # We return a list, because this will get added to the existing list
184
+ return {"messages": [response]}
185
+
186
+ @trace_agent("rewrite")
187
+ def rewrite(state):
188
+ """
189
+ Transform the query to produce a better question.
190
+
191
+ Args:
192
+ state (messages): The current state
193
+
194
+ Returns:
195
+ dict: The updated state with re-phrased question
196
+ """
197
+
198
+ print("---TRANSFORM QUERY---")
199
+ messages = state["messages"]
200
+ question = messages[0].content
201
+
202
+ msg = [
203
+ HumanMessage(
204
+ content=f""" \n
205
+ Look at the input and try to reason about the underlying semantic intent / meaning. \n
206
+ Here is the initial question:
207
+ \n ------- \n
208
+ {question}
209
+ \n ------- \n
210
+ Formulate an improved question: """,
211
+ )
212
+ ]
213
+
214
+ # Grader
215
+ model = ChatOpenAI(temperature=0, model="gpt-4-0125-preview", streaming=True)
216
+ response = model.invoke(msg)
217
+ return {"messages": [response]}
218
+
219
+ @trace_tool("generate")
220
+ def generate(state):
221
+ """
222
+ Generate answer
223
+
224
+ Args:
225
+ state (messages): The current state
226
+
227
+ Returns:
228
+ dict: The updated state with re-phrased question
229
+ """
230
+ print("---GENERATE---")
231
+ messages = state["messages"]
232
+ question = messages[0].content
233
+ last_message = messages[-1]
234
+
235
+ docs = last_message.content
236
+
237
+ # Prompt
238
+ prompt = hub.pull("rlm/rag-prompt")
239
+
240
+ # LLM
241
+ llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0, streaming=True)
242
+
243
+ # Post-processing
244
+ def format_docs(docs):
245
+ return "\n\n".join(doc.page_content for doc in docs)
246
+
247
+ # Chain
248
+ rag_chain = prompt | llm | StrOutputParser()
249
+ current_span().add_metrics(
250
+ name="generate",
251
+ score=0.9,
252
+ reasoning="Generating response based on current state",
253
+ cost=0.01,
254
+ latency=0.5,
255
+ )
256
+
257
+ # Run
258
+ response = rag_chain.invoke({"context": docs, "question": question})
259
+ return {"messages": [response]}
260
+
261
+
262
+ print("*" * 20 + "Prompt[rlm/rag-prompt]" + "*" * 20)
263
+ prompt = hub.pull("rlm/rag-prompt").pretty_print() # Show what the prompt looks like
264
+
265
+ from langgraph.graph import END, StateGraph, START
266
+ from langgraph.prebuilt import ToolNode
267
+
268
+ # Define a new graph
269
+ workflow = StateGraph(AgentState)
270
+
271
+ # Define the nodes we will cycle between
272
+ workflow.add_node("agent", agent) # agent
273
+ retrieve = ToolNode([retriever_tool])
274
+ workflow.add_node("retrieve", retrieve) # retrieval
275
+ workflow.add_node("rewrite", rewrite) # Re-writing the question
276
+ workflow.add_node(
277
+ "generate", generate
278
+ ) # Generating a response after we know the documents are relevant
279
+ # Call agent node to decide to retrieve or not
280
+ workflow.add_edge(START, "agent")
281
+
282
+ # Decide whether to retrieve
283
+ workflow.add_conditional_edges(
284
+ "agent",
285
+ # Assess agent decision
286
+ tools_condition,
287
+ {
288
+ # Translate the condition outputs to nodes in our graph
289
+ "tools": "retrieve",
290
+ END: END,
291
+ },
292
+ )
293
+
294
+ # Edges taken after the `action` node is called.
295
+ workflow.add_conditional_edges(
296
+ "retrieve",
297
+ # Assess agent decision
298
+ grade_documents,
299
+ )
300
+ workflow.add_edge("generate", END)
301
+ workflow.add_edge("rewrite", "agent")
302
+
303
+ # Compile
304
+ graph = workflow.compile()
305
+
306
+ import pprint
307
+
308
+ inputs = {
309
+ "messages": [
310
+ ("user", "What does Lilian Weng say about attacks on LLMs?"),
311
+ ]
312
+ }
313
+
314
+ with tracer:
315
+ for output in graph.stream(inputs):
316
+ for key, value in output.items():
317
+ pprint.pprint(f"Output from node '{key}':")
318
+ pprint.pprint("---")
319
+ pprint.pprint(value, indent=2, width=80, depth=None)
320
+ pprint.pprint("\n---\n")