File size: 4,679 Bytes
78fb5d7
 
 
 
 
 
 
7f4c27a
78fb5d7
 
7f4c27a
78fb5d7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c0b054b
78fb5d7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f81db3e
78fb5d7
 
 
 
 
 
 
 
c0b054b
 
78fb5d7
c932f49
1ec2cb0
c932f49
78fb5d7
c932f49
aff412c
c932f49
78fb5d7
c932f49
78fb5d7
c932f49
 
 
 
78fb5d7
c932f49
78fb5d7
c932f49
 
78fb5d7
c932f49
 
 
 
 
 
 
 
78fb5d7
c932f49
 
78fb5d7
34461a3
 
 
 
 
78fb5d7
 
 
34461a3
78fb5d7
 
 
 
 
 
 
 
 
 
b378e4d
 
 
 
78fb5d7
b378e4d
 
78fb5d7
b378e4d
 
 
78fb5d7
7fb0689
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
# -*- coding: utf-8 -*-
# Imports
import asyncio
import os
import openai

from typing import List, Optional
# from pydantic import BaseModel, Field

# from langchain.prompts import ChatPromptTemplate
# from langchain.pydantic_v1 import BaseModel
# from langchain.utils.openai_functions import convert_pydantic_to_openai_function
from llama_index.tools import FunctionTool
from llama_index.vector_stores.types import (
    VectorStoreInfo,
    MetadataInfo,
    ExactMatchFilter,
    MetadataFilters,
)
from llama_index.agent import OpenAIAgent
from llama_index.retrievers import VectorIndexRetriever
from llama_index.query_engine import RetrieverQueryEngine

from typing import List, Tuple, Any
from pydantic import BaseModel, Field
from llama_index import load_index_from_storage
from llama_index import set_global_handler
import llama_index
from llama_index.embeddings import OpenAIEmbedding
from llama_index import ServiceContext
from llama_index.llms import OpenAI
from llama_index import GPTVectorStoreIndex

set_global_handler("wandb", run_args={"project": "final-project-v1"})
wandb_callback = llama_index.global_handler

from dotenv import load_dotenv
load_dotenv()

openai.api_key = os.environ['OPENAI_API_KEY']

top_k = 3

vector_store_info = VectorStoreInfo(
    content_info="transcripts of earnings calls",
    metadata_info=[MetadataInfo(
        name="title",
        type="str",
        description="Title of the earnings call",
    ),
    MetadataInfo(
        name="period",
        type="str",
        description="Period of the earnings call"
    ),
    MetadataInfo(
        name="ticker",
        type="str",
        description="Ticker of the company"
    ),
    MetadataInfo(
        name="year",
        type="str",
        description="Year of the earnings call"
    ),
    MetadataInfo(
        name="quarter",
        type="str",
        description="Quarter of the earnings call"
    ),
    MetadataInfo(
        name="path",
        type="str",
        description="Path to the earnings call"
    ),
])

class AutoRetrieveModel(BaseModel):
    query: str = Field(..., description="natural language query string")
    filter_key_list: List[str] = Field(
        ..., description="List of metadata filter field names"
    )
    filter_value_list: List[str] = Field(
        ...,
        description=(
            "List of metadata filter field values (corresponding to names specified in filter_key_list)"
        )
    )

embed_model = OpenAIEmbedding()
chunk_size = 500

llm = OpenAI(
    temperature=0,
    model="gpt-4" ### YOUR CODE HERE
)

service_context = ServiceContext.from_defaults(
    llm=llm,
    chunk_size=chunk_size,
    embed_model=embed_model,
)

index = GPTVectorStoreIndex.from_documents([], service_context=service_context)


# Main function to extract information
async def extract_information():
    # Make sure to use a recent model that supports tools

    storage_context = wandb_callback.load_storage_context(
        artifact_url="llmop/final-project-v1/earnings-index:v1"
    )

    index = load_index_from_storage(storage_context, service_context=service_context)

    def auto_retrieve_fn(
        query: str, filter_key_list: List[str], filter_value_list: List[str]
    ):
        """Auto retrieval function.

        Performs auto-retrieval from a vector database, and then applies a set of filters.

        """
        query = query or "Query"

        exact_match_filters = [
            ExactMatchFilter(key=k, value=v)
            for k, v in zip(filter_key_list, filter_value_list)
        ]
        retriever = VectorIndexRetriever(
            index, filters=MetadataFilters(filters=exact_match_filters), top_k=top_k
        )
        query_engine = RetrieverQueryEngine.from_args(retriever, service_context=service_context)

        response = query_engine.query(query)
        return str(response)
    
    description = f"""
        Who is the CEO of MSFT
        The vector database schema is given below:
        {vector_store_info.json()}
    """
    auto_retrieve_tool = FunctionTool.from_defaults(
        fn=auto_retrieve_fn,
        name="earnings-transcripts",
        description=description,
        fn_schema=AutoRetrieveModel
    )

    agent = OpenAIAgent.from_tools(
        tools=[auto_retrieve_tool],
    )

    return agent


# if __name__ == "__main__":
#     text = "Who is the CEO of MSFT."
#     chain = extract_information()
#     print(str(chain.chat(text)))

#     async def extract_information_async(message: str):
#         return str(chain.chat(text))

#     async def main():
#         res = await extract_information_async(text)
#         print(res)

    # asyncio.run(main())