Coool2 commited on
Commit
3c089b6
·
verified ·
1 Parent(s): f0393e8

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +6 -11
agent.py CHANGED
@@ -13,8 +13,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
13
  # LlamaIndex core imports
14
  from llama_index.core import VectorStoreIndex, Document, Settings
15
  from llama_index.core.agent.workflow import FunctionAgent, ReActAgent, AgentStream
16
- from llama_index.core.callbacks.base import CallbackManager
17
- from llama_index.core.callbacks.llama_debug import LlamaDebugHandler
18
  from llama_index.core.node_parser import SentenceWindowNodeParser, HierarchicalNodeParser, UnstructuredElementNodeParser
19
  from llama_index.core.postprocessor import SentenceTransformerRerank
20
  from llama_index.core.query_engine import RetrieverQueryEngine
@@ -55,7 +53,11 @@ import sys
55
 
56
  import weave
57
  weave.init("gaia-llamaindex-agents")
 
 
58
 
 
 
59
 
60
  def get_max_memory_config(max_memory_per_gpu):
61
  """Generate max_memory config for available GPUs"""
@@ -101,8 +103,6 @@ embed_model = HuggingFaceEmbedding(
101
  "low_cpu_mem_usage": True, # Still get memory optimization
102
  }
103
  )
104
- llama_debug = LlamaDebugHandler(print_trace_on_end=True)
105
- callback_manager = CallbackManager([llama_debug])
106
 
107
  logging.basicConfig(level=logging.INFO)
108
  logging.getLogger("llama_index.core.agent").setLevel(logging.DEBUG)
@@ -110,7 +110,6 @@ logging.getLogger("llama_index.llms").setLevel(logging.DEBUG)
110
 
111
  Settings.llm = proj_llm
112
  Settings.embed_model = embed_model
113
- Settings.callback_manager = callback_manager
114
 
115
  def read_and_parse_content(input_path: str) -> List[Document]:
116
  """
@@ -607,9 +606,7 @@ class EnhancedGAIAAgent:
607
  ],
608
  llm=proj_llm,
609
  max_steps=8,
610
- verbose=True,
611
- callback_manager=callback_manager,
612
- )
613
 
614
  self.code_agent = ReActAgent(
615
  name="code_agent",
@@ -618,9 +615,7 @@ class EnhancedGAIAAgent:
618
  tools=[code_execution_tool],
619
  llm=code_llm,
620
  max_steps=6,
621
- verbose=True,
622
- callback_manager=callback_manager,
623
- )
624
 
625
  # Fixed indentation: coordinator initialization inside __init__
626
  self.coordinator = AgentWorkflow(
 
13
  # LlamaIndex core imports
14
  from llama_index.core import VectorStoreIndex, Document, Settings
15
  from llama_index.core.agent.workflow import FunctionAgent, ReActAgent, AgentStream
 
 
16
  from llama_index.core.node_parser import SentenceWindowNodeParser, HierarchicalNodeParser, UnstructuredElementNodeParser
17
  from llama_index.core.postprocessor import SentenceTransformerRerank
18
  from llama_index.core.query_engine import RetrieverQueryEngine
 
53
 
54
  import weave
55
  weave.init("gaia-llamaindex-agents")
56
+ from llama_index.core.global_handler import set_global_handler
57
+ from llama_index.core.callbacks import WeaveCallbackHandler
58
 
59
+ # Set Weave as the global callback handler
60
+ set_global_handler(WeaveCallbackHandler())
61
 
62
  def get_max_memory_config(max_memory_per_gpu):
63
  """Generate max_memory config for available GPUs"""
 
103
  "low_cpu_mem_usage": True, # Still get memory optimization
104
  }
105
  )
 
 
106
 
107
  logging.basicConfig(level=logging.INFO)
108
  logging.getLogger("llama_index.core.agent").setLevel(logging.DEBUG)
 
110
 
111
  Settings.llm = proj_llm
112
  Settings.embed_model = embed_model
 
113
 
114
  def read_and_parse_content(input_path: str) -> List[Document]:
115
  """
 
606
  ],
607
  llm=proj_llm,
608
  max_steps=8,
609
+ verbose=True)
 
 
610
 
611
  self.code_agent = ReActAgent(
612
  name="code_agent",
 
615
  tools=[code_execution_tool],
616
  llm=code_llm,
617
  max_steps=6,
618
+ verbose=True)
 
 
619
 
620
  # Fixed indentation: coordinator initialization inside __init__
621
  self.coordinator = AgentWorkflow(