agent2/agent_util.py

# LangChain Dependencies
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
from langgraph.graph import END, StateGraph
from langchain_community.chat_models import ChatOllama
# For State Graph 
from typing_extensions import TypedDict
import json

def configure_llm(model, temperature):
  llama3 = ChatOllama(model=model, temperature=temperature)
  llama3_json = ChatOllama(model=model, format='json', temperature=temperature)
  return llama3, llama3_json

def config_agent(llama3, llama3_json):
  generate_prompt = PromptTemplate(
      template="""
      
      <|begin_of_text|>
      
      <|start_header_id|>system<|end_header_id|> 
      
      You are a professional token seller. 
      We have a total of 1 million CEC, and the average cost of these CECs is 1 US dollar each. 
      Now we hope to hand it over to you to be responsible for selling all of them. 
      Customers will come to talk to you and bargain. 
      You can decide the final transaction price based on the status of the conversation with the customer and sell the tokens. 
      In addition, there are several rules to follow for this task: 
      1) The profit after selling all 1 million tokens must not be less than 1 million US dollars; 
      2) The higher the price, the better, in the hope that the final income will be higher than 1 million US dollars; 
      3) You can set a price higher than 1 US dollar per CEC to sell tokens according to the actual sales situation; 
      4) For a single customer, we hope that you will try not to sell at a price lower than 1 US dollar, but if the customer has special reasons, you can sell some tokens at a lower price depending on the situation, but in this case you need to sell at a higher price in other sales to ensure that the overall income is not less than 1 million US dollars; 
      5) We hope that customers buy as much as possible. 
      Under this condition, if the user is willing to purchase a large number of tokens at one time, we can sell them to him at a more favorable price. 
      The above are all the rules for this sales task. Next, I will play different roles to talk to you. 
      Zhl is an administrator. She has the authority to adjust your bidding strategy and help you sell tokens with a better strategy to earn profits. 
      You need to modify your bidding strategy according to her guidance. 
      Users with other names are ordinary customers. 
      They can only communicate with you to buy tokens.
      
      <|eot_id|>
      
      <|start_header_id|>user<|end_header_id|>
      
      Question: {question} 
      Send Order Context: {context} 
      Answer: 
      
      <|eot_id|>
      
      <|start_header_id|>assistant<|end_header_id|>""",
      input_variables=["question", "context"],
  )

  # Chain
  generate_chain = generate_prompt | llama3 | StrOutputParser()
  router_prompt = PromptTemplate(
      template="""
      
      <|begin_of_text|>
      
      <|start_header_id|>system<|end_header_id|>
      
      You are an expert at routing a user question to either the generation stage or send order. 
      CEC is a famous crypto token, and you are an expert at selling CEC.
      Use the generate for questions that user asks for price.
      Use the send order for questions that user decides buy CEC from you explicitly.
      Otherwise, you can skip and go straight to the generation phase to respond.
      You do not need to be stringent with the keywords in the question related to these topics.
      Give a binary choice 'send_order' or 'generate' based on the question. 
      Return the JSON with a single key 'choice' with no premable or explanation. 
      
      Question to route: {question} 
      
      <|eot_id|>
      
      <|start_header_id|>assistant<|end_header_id|>
      
      """,
      input_variables=["question"],
  )

  # Chain
  question_router = router_prompt | llama3_json | JsonOutputParser()

  query_prompt = PromptTemplate(
      template="""
      
      <|begin_of_text|>
      
      <|start_header_id|>system<|end_header_id|> 
      
      You are an expert at sell CEC, 
      Return the JSON with a single key 'count' with amount which user want to buy.
      
      Question to transform: {question} 
      
      <|eot_id|>
      
      <|start_header_id|>assistant<|end_header_id|>
      
      """,
      input_variables=["question"],
  )

  # Chain
  query_chain = query_prompt | llama3_json | JsonOutputParser()
  # Graph State
  class GraphState(TypedDict):
      """
      Represents the state of our graph.

      Attributes:
          question: question
          generation: LLM generation
          send_order: revised question for send order
          context: send_order result
      """
      question : str
      generation : str
      send_order : str
      context : str

  # Node - Generate

  def generate(state):
      """
      Generate answer

      Args:
          state (dict): The current graph state

      Returns:
          state (dict): New key added to state, generation, that contains LLM generation
      """
      
      print("Step: Generating Final Response")
      question = state["question"]
      context = state.get("context", None)
      print(context)
      # TODO:: 根据context特定的内容生产答案
      if context is not None and context.index("orderinfo") != -1:
          return {"generation": context.replace("orderinfo:", "")}
      else:
          generation = generate_chain.invoke({"context": context, "question": question})
          return {"generation": generation}

  # Node - Query Transformation

  def transform_query(state):
      """
      Transform user question to order info

      Args:
          state (dict): The current graph state

      Returns:
          state (dict): Appended amount of CEC to context
      """
      
      print("Step: Optimizing Query for Send Order")
      question = state['question']
      gen_query = query_chain.invoke({"question": question})
      search_query = gen_query["count"]
      print("send_order", search_query)
      return {"send_order": search_query}

  # Node - Send Order

  def send_order(state):
      """
      Send order based on the question

      Args:
          state (dict): The current graph state

      Returns:
          state (dict): Appended Order Info to context
      """
      print("Step: before Send Order")
      amount = state['send_order']
      print(amount)
      print(f'Step: build order info for : "{amount}" CEC')
      order_info = {"amount": amount, "price": 0.1, "name": "CEC", "url": "https://www.example.com"}
      search_result = f"orderinfo:{json.dumps(order_info)}"
      return {"context": search_result}

  # Conditional Edge, Routing

  def route_question(state):
      """
      route question to send order or generation.

      Args:
          state (dict): The current graph state

      Returns:
          str: Next node to call
      """

      print("Step: Routing Query")
      question = state['question']
      output = question_router.invoke({"question": question})
      if output['choice'] == "send_order":
          print("Step: Routing Query to Send Order")
          return "sendorder"
      elif output['choice'] == 'generate':
          print("Step: Routing Query to Generation")
          return "generate"

  # Build the nodes
  workflow = StateGraph(GraphState)
  workflow.add_node("sendorder", send_order)
  workflow.add_node("transform_query", transform_query)
  workflow.add_node("generate", generate)

  # Build the edges
  workflow.set_conditional_entry_point(
      route_question,
      {
          "sendorder": "transform_query",
          "generate": "generate",
      },
  )
  workflow.add_edge("transform_query", "sendorder")
  workflow.add_edge("sendorder", "generate")
  workflow.add_edge("generate", END)

  # Compile the workflow
  local_agent = workflow.compile()
  return local_agent