igrisds commited on
Commit
48ee539
1 Parent(s): 0428aad

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -4
app.py CHANGED
@@ -29,9 +29,9 @@ def create_bedrock_client():
29
  aws_session_token=AWS_SESSION
30
  )
31
 
 
32
  def create_s3_client():
33
 
34
- # Create an S3 client
35
  return boto3.client(
36
  's3',
37
  aws_access_key_id=AWS_ACCESS_KEY_ID,
@@ -39,6 +39,7 @@ def create_s3_client():
39
  aws_session_token=AWS_SESSION
40
  )
41
 
 
42
  def read_json_from_s3():
43
 
44
  response = s3_client.get_object(Bucket=BUCKET_NAME, Key=EXTRACTIONS_PATH)
@@ -47,6 +48,7 @@ def read_json_from_s3():
47
 
48
  return json_content
49
 
 
50
  def get_titan_embedding(bedrock_client, doc_name, text, attempt=0, cutoff=10000):
51
  """
52
  Retrieves a text embedding for a given document using the Amazon Titan Embedding model.
@@ -87,9 +89,9 @@ def get_titan_embedding(bedrock_client, doc_name, text, attempt=0, cutoff=10000)
87
  response_body = json.loads(response['body'].read())
88
  return response_body.get('embedding')
89
 
 
90
  def ask_ds(message, history):
91
 
92
-
93
  if len(message) == 0:
94
  return
95
 
@@ -110,7 +112,6 @@ def ask_ds(message, history):
110
  for file, _ in top_3:
111
  similar_content += extractions[file]['content'] + '\n'
112
 
113
-
114
  # Invoke
115
  response = bedrock_client.invoke_model_with_response_stream(
116
  modelId="anthropic.claude-3-sonnet-20240229-v1:0",
@@ -145,6 +146,7 @@ def ask_ds(message, history):
145
  ),
146
  )
147
 
 
148
  all_text = ''
149
  stream = response.get('body')
150
  if stream:
@@ -159,13 +161,15 @@ def ask_ds(message, history):
159
  except:
160
  pass
161
 
 
162
  output = '\n\nCheck out the following documents for more information:\n'
163
  for file, sim in top_3:
164
- output += f"\n{file.replace('.txt', '.pdf')} ({(1-sim) * 100}%)"
165
 
166
  yield all_text + output
167
 
168
 
 
169
  bedrock_client = create_bedrock_client()
170
  s3_client = create_s3_client()
171
  extractions = read_json_from_s3()
 
29
  aws_session_token=AWS_SESSION
30
  )
31
 
32
+ # Create AWS S3 client using environment variables
33
  def create_s3_client():
34
 
 
35
  return boto3.client(
36
  's3',
37
  aws_access_key_id=AWS_ACCESS_KEY_ID,
 
39
  aws_session_token=AWS_SESSION
40
  )
41
 
42
+ # Read JSON directly into mem from S3
43
  def read_json_from_s3():
44
 
45
  response = s3_client.get_object(Bucket=BUCKET_NAME, Key=EXTRACTIONS_PATH)
 
48
 
49
  return json_content
50
 
51
+ # Get AWS Titan embedding of text
52
  def get_titan_embedding(bedrock_client, doc_name, text, attempt=0, cutoff=10000):
53
  """
54
  Retrieves a text embedding for a given document using the Amazon Titan Embedding model.
 
89
  response_body = json.loads(response['body'].read())
90
  return response_body.get('embedding')
91
 
92
+ # Main Chat
93
  def ask_ds(message, history):
94
 
 
95
  if len(message) == 0:
96
  return
97
 
 
112
  for file, _ in top_3:
113
  similar_content += extractions[file]['content'] + '\n'
114
 
 
115
  # Invoke
116
  response = bedrock_client.invoke_model_with_response_stream(
117
  modelId="anthropic.claude-3-sonnet-20240229-v1:0",
 
146
  ),
147
  )
148
 
149
+ # Stream the response
150
  all_text = ''
151
  stream = response.get('body')
152
  if stream:
 
161
  except:
162
  pass
163
 
164
+ # Print relevant files
165
  output = '\n\nCheck out the following documents for more information:\n'
166
  for file, sim in top_3:
167
+ output += f"\n{file.replace('.txt', '.pdf')}"
168
 
169
  yield all_text + output
170
 
171
 
172
+ # Create necessary services and collect data
173
  bedrock_client = create_bedrock_client()
174
  s3_client = create_s3_client()
175
  extractions = read_json_from_s3()