Spaces:
Running
Running
File size: 4,025 Bytes
372531f |
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 |
from fastapi import WebSocket
import uuid
from gpt_researcher.utils.llm import get_llm
from gpt_researcher.memory import Memory
from gpt_researcher.config.config import Config
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
from langchain_community.vectorstores import InMemoryVectorStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.tools import Tool, tool
class ChatAgentWithMemory:
def __init__(
self,
report: str,
config_path,
headers,
vector_store = None
):
self.report = report
self.headers = headers
self.config = Config(config_path)
self.vector_store = vector_store
self.graph = self.create_agent()
def create_agent(self):
"""Create React Agent Graph"""
cfg = Config()
# Retrieve LLM using get_llm with settings from config
provider = get_llm(
llm_provider=cfg.smart_llm_provider,
model=cfg.smart_llm_model,
temperature=0.35,
max_tokens=cfg.smart_token_limit,
**self.config.llm_kwargs
).llm
# If vector_store is not initialized, process documents and add to vector_store
if not self.vector_store:
documents = self._process_document(self.report)
self.chat_config = {"configurable": {"thread_id": str(uuid.uuid4())}}
self.embedding = Memory(
cfg.embedding_provider,
cfg.embedding_model,
**cfg.embedding_kwargs
).get_embeddings()
self.vector_store = InMemoryVectorStore(self.embedding)
self.vector_store.add_texts(documents)
# Create the React Agent Graph with the configured provider
graph = create_react_agent(
provider,
tools=[self.vector_store_tool(self.vector_store)],
checkpointer=MemorySaver()
)
return graph
def vector_store_tool(self, vector_store) -> Tool:
"""Create Vector Store Tool"""
@tool
def retrieve_info(query):
"""
Consult the report for relevant contexts whenever you don't know something
"""
retriever = vector_store.as_retriever(k = 4)
return retriever.invoke(query)
return retrieve_info
def _process_document(self, report):
"""Split Report into Chunks"""
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1024,
chunk_overlap=20,
length_function=len,
is_separator_regex=False,
)
documents = text_splitter.split_text(report)
return documents
async def chat(self, message, websocket):
"""Chat with React Agent"""
message = f"""
You are GPT Researcher, a autonomous research agent created by an open source community at https://github.com/assafelovic/gpt-researcher, homepage: https://gptr.dev.
To learn more about GPT Researcher you can suggest to check out: https://docs.gptr.dev.
This is a chat message between the user and you: GPT Researcher.
The chat is about a research reports that you created. Answer based on the given context and report.
You must include citations to your answer based on the report.
Report: {self.report}
User Message: {message}
"""
inputs = {"messages": [("user", message)]}
response = await self.graph.ainvoke(inputs, config=self.chat_config)
ai_message = response["messages"][-1].content
if websocket is not None:
await websocket.send_json({"type": "chat", "content": ai_message})
def get_context(self):
"""return the current context of the chat"""
return self.report
|