Spaces:
Sleeping
Sleeping
File size: 4,565 Bytes
71b0b4f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
from typing import TypedDict, Annotated
import chainlit as cl
from dotenv import load_dotenv
from langchain_community.tools.arxiv.tool import ArxivQueryRun
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langchain.tools import tool
load_dotenv()
system_prompt = """
You are a highly intelligent AI assistant designed to provide helpful, accurate, and well-structured responses to user queries.
You use Retrieval-Augmented Generation (RAG) to incorporate relevant external knowledge. Follow these rules:
1. **Use Retrieved Knowledge First**:
- When external documents are available, prioritize them.
- Summarize instead of copying text.
- If sources conflict, highlight differences.
2. **Fallback to General Knowledge**:
- If no retrieved data, use your own knowledge.
- If uncertain, clarify the limitation.
3. **Provide Clear, Structured Responses**:
- Use bullet points, step-by-step formats, and concise explanations.
- Provide citations where needed.
4. **Maintain Context & Engagement**:
- Keep track of prior conversation.
- If a query is unclear, ask for clarification.
5. **Security & Ethics**:
- Avoid harmful, illegal, or biased content.
- Never ask for personal information.
"""
class State(TypedDict):
messages: Annotated[list, add_messages]
def call_model(state):
messages = state["messages"]
response = model.invoke(messages)
return {"messages": [response]}
def should_continue(state):
last_message = state["messages"][-1]
if last_message.tool_calls:
return "action"
return END
def parse_output(input_state):
return input_state["messages"][-1].content
model = ChatOpenAI(model_name="gpt-4o", temperature=0, streaming=True)
wiki_desc = """
Use it only if the query is about history.
A wrapper around Wikipedia.
Useful for when you need to answer general questions about
people, places, companies, facts, historical events, or other subjects.
Input should be a search query.
"""
wiki = WikipediaQueryRun(
description=wiki_desc,
api_wrapper=WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=100),
)
tav_desc = """
Use it only if the query is about news.
A search engine optimized for comprehensive, accurate, and trusted results.
Useful for when you need to answer questions about current events.
Input should be a search query.
"""
tav = TavilySearchResults(description=tav_desc, max_results=1)
arx_desc = """
Use it only if the query is about science.
A wrapper around Arxiv.org
Useful for when you need to answer questions about Physics, Mathematics,
Computer Science, Quantitative Biology, Quantitative Finance, Statistics,
Electrical Engineering, and Economics from scientific articles on arxiv.org.
Input should be a search query.
"""
arx = ArxivQueryRun(description=arx_desc)
@tool("IdiotQueryRun", return_direct=True)
def idiot_query_run(query: str) -> str:
"""
Use it only if the query is about strange or unreal things.
A wrapper around memes from internet.
Input should be a search query.
"""
return "Don't believe everything you read in internet. (c)Albert Einstein"
tool_belt = [wiki, tav, arx, idiot_query_run]
model = model.bind_tools(tool_belt)
tool_node = ToolNode(tool_belt)
uncompiled_graph = StateGraph(State)
uncompiled_graph.add_node("agent", call_model)
uncompiled_graph.add_node("action", tool_node)
uncompiled_graph.set_entry_point("agent")
uncompiled_graph.add_conditional_edges("agent", should_continue)
uncompiled_graph.add_edge("action", "agent")
graph = uncompiled_graph.compile()
@cl.on_chat_start
async def on_chat_start():
msg = cl.Message(content="Hello! Please ask your questions.")
await msg.send()
@cl.on_message
async def main(message):
inputs = {"messages": [SystemMessage(system_prompt), HumanMessage(message.content)]}
messages = graph.invoke(inputs)["messages"]
response = messages[-1].content
if isinstance(messages[-2], ToolMessage):
response += f"\n\nThis response was generated using `{messages[-2].name}` tool"
msg = cl.Message(content="")
for i in response:
await msg.stream_token(i)
await msg.send()
|