File size: 3,656 Bytes
03fe216
fe5564b
 
 
 
 
 
03fe216
5286c9a
fe5564b
5286c9a
fe5564b
5286c9a
fe5564b
 
 
 
 
 
5286c9a
fe5564b
 
 
 
 
 
 
 
 
5286c9a
fe5564b
 
 
 
 
 
5286c9a
fe5564b
 
03fe216
fe5564b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
03fe216
fe5564b
03fe216
fe5564b
03fe216
fe5564b
 
5286c9a
fe5564b
 
 
5286c9a
fe5564b
 
 
 
 
 
 
 
 
03fe216
5286c9a
fe5564b
5286c9a
fe5564b
5286c9a
 
 
fe5564b
5286c9a
 
fe5564b
5286c9a
 
 
 
 
 
 
 
 
66c3a38
5286c9a
 
 
 
 
 
 
fe5564b
5286c9a
 
fe5564b
5286c9a
03fe216
5286c9a
 
 
 
 
 
 
 
303aa50
5286c9a
03fe216
 
5286c9a
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
import gradio as gr
import os
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.prompts import PromptTemplate
from langchain_community.vectorstores import Chroma

def create_qa_chain():
    """
    Create the QA chain with the loaded vectorstore
    """
    # Initialize embeddings and load vectorstore
    embeddings = OpenAIEmbeddings()
    vectorstore = Chroma(
        persist_directory="./vectorstore",
        embedding_function=embeddings
    )
    
    # Set up retriever
    retriever = vectorstore.as_retriever(
        search_type="mmr",
        search_kwargs={
            "k": 6,
            "fetch_k": 20,
            "lambda_mult": 0.3,
        }
    )
    
    # Set up memory
    memory = ConversationBufferMemory(
        memory_key="chat_history",
        return_messages=True,
        output_key='answer'
    )
    
    # Create prompt template
    qa_prompt = PromptTemplate.from_template("""You are an expert technical writer specializing in API documentation. 
When describing API endpoints, structure your response in this exact format:

1. Start with the HTTP method and base URI structure
2. List all key parameters with:
   - Parameter name in bold (**parameter**)
   - Type and requirement status
   - Clear description
   - Example values where applicable
3. Show complete example requests with:
   - Basic example
   - Full example with all parameters
   - Headers included
4. Include any relevant response information
   
Use markdown formatting for:
- Code blocks with syntax highlighting
- Bold text for important terms
- Clear section separation

Context: {context}

Question: {question}

Technical answer (following the exact structure above):""")
    
    # Create the chain
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm=ChatOpenAI(
            temperature=0.1,
            model_name="gpt-4-turbo-preview"
        ),
        retriever=retriever,
        memory=memory,
        return_source_documents=True,
        combine_docs_chain_kwargs={"prompt": qa_prompt},
        verbose=False
    )
    
    return qa_chain

def chat(message, history):
    """
    Process chat messages and return responses
    """
    # Get or create QA chain
    if not hasattr(chat, 'qa_chain'):
        chat.qa_chain = create_qa_chain()
    
    # Get response
    result = chat.qa_chain({"question": message})
    
    # Format sources
    sources = "\n\nSources:\n"
    seen_components = set()
    shown_sources = 0
    
    for doc in result["source_documents"]:
        component = doc.metadata.get('component', '')
        title = doc.metadata.get('title', '')
        combo = (component, title)
        
        if combo not in seen_components and shown_sources < 3:
            seen_components.add(combo)
            shown_sources += 1
            sources += f"\nSource {shown_sources}:\n"
            sources += f"Title: {title}\n"
            sources += f"Component: {component}\n"
            sources += f"Content: {doc.page_content[:300]}...\n"
    
    # Combine response with sources
    full_response = result["answer"] + sources
    
    return full_response

demo = gr.ChatInterface(
    chat,
    title="Apple Music API Documentation Assistant",
    description="Ask questions about the Apple Music API documentation.",
    examples=[
        "How to search for songs on Apple Music API?",
        "What are the required parameters for searching songs?",
        "Show me an example request with all parameters"
    ]
)

if __name__ == "__main__":
    demo.launch()