|
| 1 | +"""LlamaIndex zero-code OTLP example.""" |
| 2 | +import asyncio |
| 3 | +import os |
| 4 | +import random |
| 5 | + |
| 6 | +from dotenv import load_dotenv |
| 7 | +from llama_index.core.agent.workflow import FunctionAgent |
| 8 | +from llama_index.core.tools import FunctionTool |
| 9 | +from llama_index.llms.openai_like import OpenAILike |
| 10 | +from opentelemetry import trace |
| 11 | +from opentelemetry._logs import set_logger_provider |
| 12 | +from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter |
| 13 | +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter |
| 14 | +from opentelemetry.instrumentation.openai_v2 import OpenAIInstrumentor |
| 15 | +from opentelemetry.sdk._logs import LoggerProvider |
| 16 | +from opentelemetry.sdk._logs.export import BatchLogRecordProcessor |
| 17 | +from opentelemetry.sdk.resources import Resource |
| 18 | +from opentelemetry.sdk.trace import TracerProvider |
| 19 | +from opentelemetry.sdk.trace.export import BatchSpanProcessor |
| 20 | + |
| 21 | +load_dotenv(override=True) |
| 22 | + |
| 23 | + |
| 24 | +def roll_die(sides: int) -> int: |
| 25 | + """Roll a die and return the result.""" |
| 26 | + return random.randint(1, sides) |
| 27 | + |
| 28 | + |
| 29 | +def check_prime(number: int) -> bool: |
| 30 | + """Check if a number is prime.""" |
| 31 | + return number >= 2 and all(number % i for i in range(2, int(number**0.5) + 1)) |
| 32 | + |
| 33 | + |
| 34 | +async def main(): |
| 35 | + if not os.getenv("OPENAI_API_KEY"): |
| 36 | + print("OPENAI_API_KEY not set.") |
| 37 | + return |
| 38 | + |
| 39 | + os.environ["OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"] = "true" |
| 40 | + os.environ.setdefault("OTEL_RESOURCE_ATTRIBUTES", |
| 41 | + "agentevals.eval_set_id=llama_index_eval,agentevals.session_name=llama-index-zero-code") |
| 42 | + |
| 43 | + resource = Resource.create() |
| 44 | + |
| 45 | + tp = TracerProvider(resource=resource) |
| 46 | + tp.add_span_processor(BatchSpanProcessor(OTLPSpanExporter(), schedule_delay_millis=1000)) |
| 47 | + trace.set_tracer_provider(tp) |
| 48 | + |
| 49 | + lp = LoggerProvider(resource=resource) |
| 50 | + lp.add_log_record_processor(BatchLogRecordProcessor(OTLPLogExporter(), schedule_delay_millis=1000)) |
| 51 | + set_logger_provider(lp) |
| 52 | + |
| 53 | + OpenAIInstrumentor().instrument() |
| 54 | + |
| 55 | + llm = OpenAILike( |
| 56 | + model=os.environ.get("OPENAI_MODEL", "gpt-4o-mini"), |
| 57 | + api_base=os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1"), |
| 58 | + is_chat_model=True, is_function_calling_model=True, |
| 59 | + ) |
| 60 | + agent = FunctionAgent( |
| 61 | + tools=[FunctionTool.from_defaults(fn=roll_die), FunctionTool.from_defaults(fn=check_prime)], |
| 62 | + llm=llm, |
| 63 | + system_prompt="Use roll_die for dice rolls. Use check_prime to check if a number is prime.", |
| 64 | + ) |
| 65 | + |
| 66 | + queries = ["Hi! Can you help me?", "Roll a 20-sided die.", "Roll a 6-sided die and check if the result is prime."] |
| 67 | + for i, query in enumerate(queries, 1): |
| 68 | + print(f"\n[{i}/{len(queries)}] User: {query}") |
| 69 | + result = await agent.run(query) |
| 70 | + print(f" Response: {result.response.content}") |
| 71 | + |
| 72 | + print() |
| 73 | + tp.force_flush() |
| 74 | + lp.force_flush() |
| 75 | + print("All traces flushed.") |
| 76 | + |
| 77 | + |
| 78 | +if __name__ == "__main__": |
| 79 | + asyncio.run(main()) |
0 commit comments