1+ import os
2+ from typing import Dict , List
3+ from dotenv import load_dotenv
4+ from fastapi import FastAPI , HTTPException
5+ from pydantic import BaseModel
6+ from langchain_aws import ChatBedrock
7+ from langchain .prompts import ChatPromptTemplate
8+ from langchain .chains import LLMChain
9+ from opentelemetry import trace
10+ from opentelemetry .sdk .trace import TracerProvider
11+ from opentelemetry .sdk .trace .export import BatchSpanProcessor , ConsoleSpanExporter
12+ from opentelemetry .exporter .otlp .proto .http .trace_exporter import OTLPSpanExporter
13+ from openinference .instrumentation .langchain import LangChainInstrumentor
14+
15+ # Load environment variables
16+ load_dotenv ()
17+
18+ # Set up OpenTelemetry with BOTH exporters
19+ tracer_provider = TracerProvider ()
20+
21+ # Add Console exporter
22+ console_exporter = ConsoleSpanExporter ()
23+ console_processor = BatchSpanProcessor (console_exporter )
24+ tracer_provider .add_span_processor (console_processor )
25+
26+ # Add OTLP exporter
27+ otlp_exporter = OTLPSpanExporter (endpoint = "http://localhost:4318/v1/traces" )
28+ otlp_processor = BatchSpanProcessor (otlp_exporter )
29+ tracer_provider .add_span_processor (otlp_processor )
30+
31+ # Set as global provider
32+ trace .set_tracer_provider (tracer_provider )
33+
34+ # Instrument LangChain with OpenInference
35+ LangChainInstrumentor ().instrument (tracer_provider = tracer_provider )
36+
37+ # Initialize FastAPI app
38+ app = FastAPI (title = "LangChain Bedrock OpenInference API" , version = "1.0.0" )
39+
40+ # Initialize the LLM with AWS Bedrock
41+ llm = ChatBedrock (
42+ model_id = "anthropic.claude-3-haiku-20240307-v1:0" ,
43+ model_kwargs = {
44+ "temperature" : 0.7 ,
45+ "max_tokens" : 500
46+ },
47+ region_name = os .getenv ("AWS_DEFAULT_REGION" , "us-west-2" )
48+ )
49+
50+ # Create a prompt template
51+ prompt = ChatPromptTemplate .from_template (
52+ "You are a helpful assistant. The user says: {input}. Provide a helpful response."
53+ )
54+
55+ # Create a chain
56+ chain = LLMChain (llm = llm , prompt = prompt )
57+
58+ # Request models
59+ class ChatRequest (BaseModel ):
60+ message : str
61+
62+ class BatchChatRequest (BaseModel ):
63+ messages : List [str ]
64+
65+ class ChatResponse (BaseModel ):
66+ response : str
67+
68+ class BatchChatResponse (BaseModel ):
69+ responses : List [Dict [str , str ]]
70+
71+ # Sample prompts for testing
72+ SAMPLE_PROMPTS = [
73+ "What is the capital of France?" ,
74+ "How do I make a cup of coffee?" ,
75+ "What are the benefits of exercise?" ,
76+ "Explain quantum computing in simple terms" ,
77+ "What's the best way to learn programming?"
78+ ]
79+
80+ @app .get ("/" )
81+ async def root ():
82+ return {
83+ "message" : "LangChain Bedrock OpenInference API is running!" ,
84+ "endpoints" : {
85+ "/chat" : "Single message chat endpoint" ,
86+ "/batch" : "Batch message processing endpoint" ,
87+ "/sample" : "Run sample prompts"
88+ }
89+ }
90+
91+ @app .post ("/chat" , response_model = ChatResponse )
92+ async def chat (request : ChatRequest ):
93+ """
94+ Chat endpoint that processes a single user message through AWS Bedrock
95+ """
96+ try :
97+ # Process the input through the chain
98+ result = await chain .ainvoke ({"input" : request .message })
99+ return ChatResponse (response = result ["text" ])
100+ except Exception as e :
101+ raise HTTPException (status_code = 500 , detail = str (e ))
102+
103+ @app .post ("/batch" , response_model = BatchChatResponse )
104+ async def batch_chat (request : BatchChatRequest ):
105+ """
106+ Batch endpoint that processes multiple messages
107+ """
108+ try :
109+ responses = []
110+ for message in request .messages :
111+ result = await chain .ainvoke ({"input" : message })
112+ responses .append ({
113+ "message" : message ,
114+ "response" : result ["text" ]
115+ })
116+ return BatchChatResponse (responses = responses )
117+ except Exception as e :
118+ raise HTTPException (status_code = 500 , detail = str (e ))
119+
120+ @app .get ("/sample" , response_model = BatchChatResponse )
121+ async def run_samples ():
122+ """
123+ Run the predefined sample prompts
124+ """
125+ try :
126+ responses = []
127+ for prompt in SAMPLE_PROMPTS :
128+ result = await chain .ainvoke ({"input" : prompt })
129+ responses .append ({
130+ "message" : prompt ,
131+ "response" : result ["text" ]
132+ })
133+ return BatchChatResponse (responses = responses )
134+ except Exception as e :
135+ raise HTTPException (status_code = 500 , detail = str (e ))
136+
137+ @app .get ("/health" )
138+ async def health ():
139+ """Health check endpoint"""
140+ return {"status" : "healthy" , "llm" : "AWS Bedrock Claude 3 Haiku" }
141+
142+ if __name__ == "__main__" :
143+ import uvicorn
144+ print ("Starting FastAPI server with AWS Bedrock and OpenInference instrumentation..." )
145+ print ("Make sure AWS credentials are configured" )
146+ print ("Server will run on http://localhost:8000" )
147+ print ("API docs available at http://localhost:8000/docs" )
148+ uvicorn .run (app , host = "0.0.0.0" , port = 8000 )
0 commit comments